diff --git a/mocket/mocket.py b/mocket/mocket.py index a2eabe39..dc893870 100644 --- a/mocket/mocket.py +++ b/mocket/mocket.py @@ -7,7 +7,6 @@ import io import collections import hashlib -import zlib import gzip from datetime import datetime, timedelta @@ -199,12 +198,10 @@ def sendall(self, data, *args, **kwargs): self.fd.seek(0) def recv(self, buffersize, flags=None): - resp = self.fd.readline(buffersize) - return resp + return self.fd.readline(buffersize) def _connect(self): # pragma: no cover if not self._connected: - print(self.true_socket) self.true_socket.connect(self._address) self._connected = True @@ -215,7 +212,9 @@ def true_sendall(self, data, *args, **kwargs): # port should be always a string port = text_type(self._port) - responses = {self._host: {port: {}}} + # prepare responses dictionary + responses = dict() + if Mocket.get_truesocket_recording_dir(): path = os.path.join( Mocket.get_truesocket_recording_dir(), @@ -224,27 +223,35 @@ def true_sendall(self, data, *args, **kwargs): # check if there's already a recorded session dumped to a JSON file try: with io.open(path) as f: - responses.update(json.load(f)) + responses = json.load(f) # if not, create a new dictionary - except (FileNotFoundError, JSONDecodeError, KeyError): + except (FileNotFoundError, JSONDecodeError): pass - response_dict = responses[self._host][port] + try: + response_dict = responses[self._host][port][req_signature] + except KeyError: + # preventing next KeyError exceptions + responses.setdefault(self._host, dict()) + responses[self._host].setdefault(port, dict()) + responses[self._host][port].setdefault(req_signature, dict()) + response_dict = responses[self._host][port][req_signature] # try to get the response from the dictionary try: - lines = response_dict[req_signature]['response'] - gzipped_lines = response_dict[req_signature]['gzip'] + lines = response_dict['response'] + gzipped_lines = response_dict['gzip'] r_lines = [] for line_no, line in enumerate(lines): line = encode_utf8(line) if line_no + 1 in gzipped_lines: gzip_buffer = io.BytesIO() gzip_file = gzip.GzipFile(mode='wb', fileobj=gzip_buffer) - gzip_file.write(line) - gzip_file.close() + try: + gzip_file.write(line) + finally: + gzip_file.close() line = gzip_buffer.getvalue() - # line = zlib.compress(line) r_lines.append(line) encoded_response = b'\r\n'.join(r_lines) written = len(encoded_response) @@ -261,9 +268,9 @@ def true_sendall(self, data, *args, **kwargs): if len(recv) < self._buflen: break - response_dict[req_signature] = dict(request=req) - lines = response_dict[req_signature]['response'] = [] - gzipped_lines = response_dict[req_signature]['gzip'] = [] + response_dict['request'] = req + lines = response_dict['response'] = [] + gzipped_lines = response_dict['gzip'] = [] # update the dictionary with the response obtained encoded_response = r.getvalue() @@ -272,7 +279,12 @@ def true_sendall(self, data, *args, **kwargs): try: line = decode_utf8(line) except UnicodeDecodeError: - line = decode_utf8(zlib.decompress(line, 16 + zlib.MAX_WBITS)) + f = gzip.GzipFile(mode='rb', fileobj=io.BytesIO(line)) + try: + line = f.read(len(line)) + finally: + f.close() + line = decode_utf8(line) gzipped_lines.append(line_no + 1) lines.append(line) @@ -280,7 +292,7 @@ def true_sendall(self, data, *args, **kwargs): # dump the resulting dictionary to a JSON file if self._truesocket_recording_dir: with io.open(path, mode='w', encoding=encoding) as f: - f.write(decode_utf8(json.dumps(responses))) + f.write(decode_utf8(json.dumps(responses, indent=4, sort_keys=True))) # write the response to the mocket socket self.fd.write(encoded_response) diff --git a/tests/main/test_http.py b/tests/main/test_http.py index 89e2fbd6..f90a0e9c 100644 --- a/tests/main/test_http.py +++ b/tests/main/test_http.py @@ -66,6 +66,23 @@ def test_truesendall_with_gzip_recording(self): assert len(responses['httpbin.org']['80'].keys()) == 1 + @mocketize(truesocket_recording_dir=recording_directory) + def test_truesendall_with_chunk_recording(self): + url = 'http://httpbin.org/range/2048?chunk_size=256' + + 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 io.open(dump_filename) as f: + responses = json.load(f) + + assert len(responses['httpbin.org']['80'].keys()) == 1 + @mocketize(truesocket_recording_dir=os.path.dirname(__file__)) def test_truesendall_with_dump_from_recording(self): requests.get('http://httpbin.org/ip', headers={"user-agent": "Fake-User-Agent"})