Addressed bug where recording requests does not save headers to file. Closes #715 (#720)

TheJacobWalters · web-flow · commit b0573d2d89b5 · 2024-06-14T09:27:31.000-07:00
diff --git a/responses/__init__.py b/responses/__init__.py
@@ -814,7 +814,11 @@ def add(
 
         if adding_headers is not None:
             kwargs.setdefault("headers", adding_headers)
-        if "content_type" in kwargs and "headers" in kwargs:
+        if (
+            "content_type" in kwargs
+            and "headers" in kwargs
+            and kwargs["headers"] is not None
+        ):
             header_keys = [header.lower() for header in kwargs["headers"]]
             if "content-type" in header_keys:
                 raise RuntimeError(
@@ -852,6 +856,7 @@ def _add_from_file(self, file_path: "Union[str, bytes, os.PathLike[Any]]") -> No
                 url=rsp["url"],
                 body=rsp["body"],
                 status=rsp["status"],
+                headers=rsp["headers"] if "headers" in rsp else None,
                 content_type=rsp["content_type"],
                 auto_calculate_content_length=rsp["auto_calculate_content_length"],
             )
diff --git a/responses/_recorder.py b/responses/_recorder.py
@@ -36,6 +36,29 @@ def _remove_nones(d: "Any") -> "Any":
     return d
 
 
+def _remove_default_headers(data: "Any") -> "Any":
+    """
+    It would be too verbose to store these headers in the file generated by the
+    record functionality.
+    """
+    if isinstance(data, dict):
+        keys_to_remove = [
+            "Content-Length",
+            "Content-Type",
+            "Date",
+            "Server",
+            "Connection",
+            "Content-Encoding",
+        ]
+        for i, response in enumerate(data["responses"]):
+            for key in keys_to_remove:
+                if key in response["response"]["headers"]:
+                    del data["responses"][i]["response"]["headers"][key]
+            if not response["response"]["headers"]:
+                del data["responses"][i]["response"]["headers"]
+    return data
+
+
 def _dump(
     registered: "List[BaseResponse]",
     destination: "Union[BinaryIO, TextIOWrapper]",
@@ -63,7 +86,8 @@ def _dump(
                 "Cannot dump response object."
                 "Probably you use custom Response object that is missing required attributes"
             ) from exc
-    dumper(_remove_nones(data), destination)
+
+    dumper(_remove_default_headers(_remove_nones(data)), destination)
 
 
 class Recorder(RequestsMock):
@@ -116,11 +140,15 @@ def _on_request(
         request.params = self._parse_request_params(request.path_url)  # type: ignore[attr-defined]
         request.req_kwargs = kwargs  # type: ignore[attr-defined]
         requests_response = _real_send(adapter, request, **kwargs)
+        headers_values = {
+            key: value for key, value in requests_response.headers.items()
+        }
         responses_response = Response(
             method=str(request.method),
             url=str(requests_response.request.url),
             status=requests_response.status_code,
             body=requests_response.text,
+            headers=headers_values,
         )
         self._registry.add(responses_response)
         return requests_response
diff --git a/responses/tests/test_recorder.py b/responses/tests/test_recorder.py
@@ -23,6 +23,7 @@ def get_data(host, port):
                 "response": {
                     "method": "GET",
                     "url": f"http://{host}:{port}/404",
+                    "headers": {"x": "foo"},
                     "body": "404 Not Found",
                     "status": 404,
                     "content_type": "text/plain",
@@ -33,6 +34,7 @@ def get_data(host, port):
                 "response": {
                     "method": "GET",
                     "url": f"http://{host}:{port}/status/wrong",
+                    "headers": {"x": "foo"},
                     "body": "Invalid status code",
                     "status": 400,
                     "content_type": "text/plain",
@@ -43,6 +45,7 @@ def get_data(host, port):
                 "response": {
                     "method": "GET",
                     "url": f"http://{host}:{port}/500",
+                    "headers": {"x": "foo"},
                     "body": "500 Internal Server Error",
                     "status": 500,
                     "content_type": "text/plain",
@@ -89,7 +92,6 @@ def run():
 
         with open(self.out_file) as file:
             data = yaml.safe_load(file)
-
         assert data == get_data(httpserver.host, httpserver.port)
 
     def test_recorder_toml(self, httpserver):
@@ -122,16 +124,27 @@ def run():
 
     def prepare_server(self, httpserver):
         httpserver.expect_request("/500").respond_with_data(
-            "500 Internal Server Error", status=500, content_type="text/plain"
+            "500 Internal Server Error",
+            status=500,
+            content_type="text/plain",
+            headers={"x": "foo"},
         )
         httpserver.expect_request("/202").respond_with_data(
-            "OK", status=202, content_type="text/plain"
+            "OK",
+            status=202,
+            content_type="text/plain",
         )
         httpserver.expect_request("/404").respond_with_data(
-            "404 Not Found", status=404, content_type="text/plain"
+            "404 Not Found",
+            status=404,
+            content_type="text/plain",
+            headers={"x": "foo"},
         )
         httpserver.expect_request("/status/wrong").respond_with_data(
-            "Invalid status code", status=400, content_type="text/plain"
+            "Invalid status code",
+            status=400,
+            content_type="text/plain",
+            headers={"x": "foo"},
         )
         url500 = httpserver.url_for("/500")
         url202 = httpserver.url_for("/202")
