respond_with_data: fix headers type hinting

csernazs · csernazs · commit 8704fb092261 · 2022-06-05T12:52:13.000+02:00
This can accept a mapping or a list of 2-element tuples. In this latter case, multiple values can be specified for the same header. A testcase is added for this. Fixes csernazs#132
diff --git a/pytest_httpserver/httpserver.py b/pytest_httpserver/httpserver.py
@@ -8,7 +8,7 @@
 from enum import Enum
 from contextlib import suppress, contextmanager
 from copy import copy
-from typing import Any, Callable, List, Mapping, MutableMapping, Optional, Tuple, Union, Pattern
+from typing import Any, Callable, Iterable, List, Mapping, MutableMapping, Optional, Tuple, Union, Pattern
 from ssl import SSLContext
 import abc
 
@@ -23,6 +23,11 @@
 URI_DEFAULT = ""
 METHOD_ALL = "__ALL"
 
+HEADERS_T = Union[
+    Mapping[str, Union[str, int, Iterable[Union[str, int]]]],
+    Iterable[Tuple[str, Union[str, int]]],
+]
+
 
 class Undefined:
     def __repr__(self):
@@ -475,7 +480,7 @@ def respond_with_data(
             self,
             response_data: Union[str, bytes] = "",
             status: int = 200,
-            headers: Optional[Mapping[str, str]] = None,
+            headers: Optional[HEADERS_T] = None,
             mimetype: Optional[str] = None,
             content_type: Optional[str] = None):
         """
diff --git a/tests/test_headers.py b/tests/test_headers.py
@@ -1,8 +1,10 @@
+import http.client
+
 import requests
 from pytest_httpserver import HTTPServer
-from werkzeug.http import parse_dict_header
-
 from pytest_httpserver.httpserver import HeaderValueMatcher
+from werkzeug.http import parse_dict_header
+from werkzeug.datastructures import Headers
 
 
 def test_custom_headers(httpserver: HTTPServer):
@@ -72,3 +74,24 @@ def test_authorization_headers(httpserver: HTTPServer):
     response = requests.get(httpserver.url_for('/'), headers=headers_with_values_in_modified_order)
     assert response.status_code == 200
     assert response.text == 'OK'
+
+
+def test_header_one_key_multiple_values(httpserver: HTTPServer):
+    httpserver.expect_request(uri="/t1").respond_with_data(headers=[("X-Foo", "123"), ("X-Foo", "456")])
+    httpserver.expect_request(uri="/t2").respond_with_data(headers={"X-Foo": ["123", "456"]})
+    httpserver.expect_request(uri="/t3").respond_with_data(headers={"X-Foo": [123, 456]})
+
+    headers = Headers()
+    headers.add("X-Foo", "123")
+    headers.add("X-Foo", "456")
+
+    httpserver.expect_request(uri="/t4").respond_with_data(headers=headers)
+
+    for uri in ("/t1", "/t2", "/t3", "/t4"):
+        conn = http.client.HTTPConnection("localhost:{}".format(httpserver.port))
+        conn.request("GET", uri)
+        response = conn.getresponse()
+        conn.close()
+
+        assert response.status == 200
+        assert response.headers.get_all("X-Foo") == ["123", "456"]
