diff --git a/mocket/async_mocket.py b/mocket/async_mocket.py index 709d225f..da3ba35a 100644 --- a/mocket/async_mocket.py +++ b/mocket/async_mocket.py @@ -7,11 +7,17 @@ async def wrapper( truesocket_recording_dir=None, strict_mode=False, strict_mode_allowed=None, + use_hex_encoding=True, *args, **kwargs, ): async with Mocketizer.factory( - test, truesocket_recording_dir, strict_mode, strict_mode_allowed, args + test, + truesocket_recording_dir, + strict_mode, + strict_mode_allowed, + use_hex_encoding, + args, ): return await test(*args, **kwargs) diff --git a/mocket/inject.py b/mocket/inject.py index 35e9da01..4eb1537f 100644 --- a/mocket/inject.py +++ b/mocket/inject.py @@ -17,6 +17,7 @@ def enable( namespace: str | None = None, truesocket_recording_dir: str | None = None, + use_hex_encoding=True, ) -> None: from mocket.mocket import Mocket from mocket.socket import ( @@ -33,6 +34,7 @@ def enable( Mocket._namespace = namespace Mocket._truesocket_recording_dir = truesocket_recording_dir + Mocket._use_hex_encoding = use_hex_encoding if truesocket_recording_dir and not os.path.isdir(truesocket_recording_dir): # JSON dumps will be saved here diff --git a/mocket/mocket.py b/mocket/mocket.py index 3476902d..bbb7feb1 100644 --- a/mocket/mocket.py +++ b/mocket/mocket.py @@ -22,6 +22,7 @@ class Mocket: _requests: ClassVar[list] = [] _namespace: ClassVar[str] = str(id(_entries)) _truesocket_recording_dir: ClassVar[str | None] = None + _use_hex_encoding: ClassVar[bool] = True enable = mocket.inject.enable disable = mocket.inject.disable @@ -96,6 +97,10 @@ def get_namespace(cls) -> str: def get_truesocket_recording_dir(cls) -> str | None: return cls._truesocket_recording_dir + @classmethod + def get_use_hex_encoding(cls) -> bool: + return cls._use_hex_encoding + @classmethod def assert_fail_if_entries_not_served(cls) -> None: """Mocket checks that all entries have been served at least once.""" diff --git a/mocket/mocketizer.py b/mocket/mocketizer.py index 2bf2b9cd..4d00805f 100644 --- a/mocket/mocketizer.py +++ b/mocket/mocketizer.py @@ -9,11 +9,13 @@ def __init__( instance=None, namespace=None, truesocket_recording_dir=None, + use_hex_encoding=True, strict_mode=False, strict_mode_allowed=None, ): self.instance = instance self.truesocket_recording_dir = truesocket_recording_dir + self.use_hex_encoding = use_hex_encoding self.namespace = namespace or str(id(self)) MocketMode().STRICT = strict_mode if strict_mode: @@ -27,6 +29,7 @@ def enter(self): Mocket.enable( namespace=self.namespace, truesocket_recording_dir=self.truesocket_recording_dir, + use_hex_encoding=self.use_hex_encoding, ) if self.instance: self.check_and_call("mocketize_setup") @@ -57,7 +60,14 @@ def check_and_call(self, method_name): method() @staticmethod - def factory(test, truesocket_recording_dir, strict_mode, strict_mode_allowed, args): + def factory( + test, + truesocket_recording_dir, + strict_mode, + strict_mode_allowed, + use_hex_encoding, + args, + ): instance = args[0] if args else None namespace = None if truesocket_recording_dir: @@ -74,6 +84,7 @@ def factory(test, truesocket_recording_dir, strict_mode, strict_mode_allowed, ar namespace=namespace, truesocket_recording_dir=truesocket_recording_dir, strict_mode=strict_mode, + use_hex_encoding=use_hex_encoding, strict_mode_allowed=strict_mode_allowed, ) @@ -83,11 +94,17 @@ def wrapper( truesocket_recording_dir=None, strict_mode=False, strict_mode_allowed=None, + use_hex_encoding=True, *args, **kwargs, ): with Mocketizer.factory( - test, truesocket_recording_dir, strict_mode, strict_mode_allowed, args + test, + truesocket_recording_dir, + strict_mode, + strict_mode_allowed, + use_hex_encoding, + args, ): return test(*args, **kwargs) diff --git a/mocket/socket.py b/mocket/socket.py index 03c6f7e5..7bcc565d 100644 --- a/mocket/socket.py +++ b/mocket/socket.py @@ -2,6 +2,7 @@ import contextlib import errno +import gzip import hashlib import json import os @@ -14,7 +15,7 @@ import urllib3.connection from typing_extensions import Self -from mocket.compat import decode_from_bytes, encode_to_bytes +from mocket.compat import decode_from_bytes, encode_to_bytes, ENCODING from mocket.entry import MocketEntry from mocket.io import MocketSocketIO from mocket.mocket import Mocket @@ -291,7 +292,20 @@ def true_sendall(self, data: ReadableBuffer, *args: Any, **kwargs: Any) -> int: # try to get the response from the dictionary try: - encoded_response = hexload(response_dict["response"]) + response = response_dict["response"] + + if Mocket.get_use_hex_encoding(): + encoded_response = hexload(response) + else: + headers, body = response.split("\r\n\r\n", 1) + + headers_bytes = headers.encode(ENCODING) + body_bytes = body.encode(ENCODING) + + if "content-encoding: gzip" in headers.lower(): + body_bytes = gzip.compress(body_bytes) + + encoded_response = headers_bytes + b"\r\n\r\n" + body_bytes # if not available, call the real sendall except KeyError: host, port = self._host, self._port @@ -316,7 +330,18 @@ def true_sendall(self, data: ReadableBuffer, *args: Any, **kwargs: Any) -> int: if Mocket.get_truesocket_recording_dir(): # update the dictionary with request and response lines response_dict["request"] = req - response_dict["response"] = hexdump(encoded_response) + + if Mocket.get_use_hex_encoding(): + response_dict["response"] = hexdump(encoded_response) + else: + headers, body = encoded_response.split(b"\r\n\r\n", 1) + + if b"content-encoding: gzip" in headers.lower(): + body = gzip.decompress(body) + + response_dict["response"] = (headers + b"\r\n\r\n" + body).decode( + ENCODING + ) with open(path, mode="w") as f: f.write( diff --git a/tests/test_http.py b/tests/test_http.py index afa31185..2133bba4 100644 --- a/tests/test_http.py +++ b/tests/test_http.py @@ -93,6 +93,58 @@ def test_truesendall_with_chunk_recording(self): assert len(responses["httpbin.local"]["80"].keys()) == 1 + def test_truesendall_with_recording_without_hex_encoding(self): + with tempfile.TemporaryDirectory() as temp_dir, Mocketizer( + truesocket_recording_dir=temp_dir, use_hex_encoding=False + ): + url = "http://httpbin.local/ip" + + requests.get(url) + resp = requests.get(url) + self.assertEqual(resp.status_code, 200) + + dump_filename = os.path.join( + Mocket.get_truesocket_recording_dir(), + Mocket.get_namespace() + ".json", + ) + with open(dump_filename) as f: + responses = json.load(f) + + for _, value in responses["httpbin.local"]["80"].items(): + self.assertIn("HTTP/1.1 200", value["response"]) + + def test_truesendall_with_gzip_recording_without_hex_encoding(self): + with tempfile.TemporaryDirectory() as temp_dir, Mocketizer( + truesocket_recording_dir=temp_dir, use_hex_encoding=False + ): + url = "http://httpbin.local/gzip" + headers = { + "Accept-Encoding": "gzip, deflate, zstd", + } + + requests.get( + url, + headers=headers, + ) + + dump_filename = os.path.join( + Mocket.get_truesocket_recording_dir(), + Mocket.get_namespace() + ".json", + ) + + with open(dump_filename) as f: + responses = json.load(f) + + for _, value in responses["httpbin.local"]["80"].items(): + self.assertIn("HTTP/1.1 200", value["response"]) + self.assertIn("gzip, deflate, zstd", value["response"]) + + resp = requests.get( + url, + headers=headers, + ) + self.assertEqual(resp.status_code, 200) + @mocketize def test_wrongpath_truesendall(self): Entry.register(