diff --git a/tests/rest/client/test_login.py b/tests/rest/client/test_login.py index 3fb77fd9dd..3a99838719 100644 --- a/tests/rest/client/test_login.py +++ b/tests/rest/client/test_login.py @@ -969,9 +969,8 @@ class CASTestCase(unittest.HomeserverTestCase): # Test that the response is HTML. self.assertEqual(channel.code, 200, channel.result) content_type_header_value = "" - for header in channel.result.get("headers", []): - if header[0] == b"Content-Type": - content_type_header_value = header[1].decode("utf8") + for header in channel.headers.getRawHeaders("Content-Type"): + content_type_header_value = header self.assertTrue(content_type_header_value.startswith("text/html")) diff --git a/tests/test_server.py b/tests/test_server.py index 0910ea5f28..45015ca653 100644 --- a/tests/test_server.py +++ b/tests/test_server.py @@ -393,7 +393,7 @@ class WrapHtmlRequestHandlerTests(unittest.TestCase): self.assertEqual(channel.code, 301) headers = channel.result["headers"] - location_headers = [v for k, v in headers if k == b"Location"] + location_headers = headers.getRawHeaders(b"Location", []) self.assertEqual(location_headers, [b"/look/an/eagle"]) def test_redirect_exception_with_cookie(self) -> None: @@ -416,9 +416,9 @@ class WrapHtmlRequestHandlerTests(unittest.TestCase): self.assertEqual(channel.code, 304) headers = channel.result["headers"] - location_headers = [v for k, v in headers if k == b"Location"] + location_headers = headers.getRawHeaders(b"Location", []) self.assertEqual(location_headers, [b"/no/over/there"]) - cookies_headers = [v for k, v in headers if k == b"Set-Cookie"] + cookies_headers = headers.getRawHeaders(b"Set-Cookie", []) self.assertEqual(cookies_headers, [b"session=yespls"]) def test_head_request(self) -> None: