Skip to content

Commit

Permalink
We should take care of gzipped lines.
Browse files Browse the repository at this point in the history
  • Loading branch information
Giorgio Salluzzo committed Feb 14, 2017
1 parent e14752b commit 949c6a5
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 3 deletions.
28 changes: 25 additions & 3 deletions mocket/mocket.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import io
import collections
import hashlib
import zlib
from datetime import datetime, timedelta

import decorator
Expand Down Expand Up @@ -216,8 +217,17 @@ def true_sendall(self, data, *args, **kwargs):

# try to get the response from the dictionary
try:
encoded_response = encode_utf8(responses[self._host][port][req_signature]['response'])
written = len(encoded_response)
written = 0
lines = responses[self._host][port][req_signature]['response']
gzip = responses[self._host][port][req_signature]['gzip']
encoded_response = byte_type()
for line_no, line in enumerate(lines):
line = encode_utf8(line)
if line_no + 1 in gzip:
line = zlib.compress(line)
written += len(line)
encoded_response += encode_utf8(line)

# if not available, call the real sendall
except KeyError:
self._connect()
Expand All @@ -230,9 +240,21 @@ def true_sendall(self, data, *args, **kwargs):
if len(recv) < self._buflen:
break

responses[self._host][port][req_signature] = dict(request=req)
lines = responses[self._host][port][req_signature]['response'] = []
gzip = responses[self._host][port][req_signature]['gzip'] = []

# update the dictionary with the response obtained
encoded_response = r.getvalue()
responses[self._host][port][req_signature] = dict(request=req, response=decode_utf8(encoded_response))

for line_no, line in enumerate(encoded_response.split(b'\r\n')):
try:
line = decode_utf8(line)
except UnicodeDecodeError:
line = decode_utf8(zlib.decompress(line, 16 + zlib.MAX_WBITS))
gzip.append(line_no + 1)

lines.append(line)

# dump the resulting dictionary to a JSON file
if self._record_truesocket:
Expand Down
15 changes: 15 additions & 0 deletions tests/main/test_http.py
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,21 @@ def test_truesendall_with_recording(self):

assert len(responses['httpbin.org']['80'].keys()) == 2

@mocketize(record_truesocket=True)
def test_truesendall_with_gzip_recording(self):
os.environ["MOCKET-RECORDING"] = tempfile.mkdtemp()
d = os.getenv("MOCKET-RECORDING")

requests.get('http://httpbin.org/gzip')
resp = requests.get('http://httpbin.org/gzip')
self.assertEqual(resp.status_code, 200)

dump_filename = os.path.join(d, Mocket.get_namespace() + '.json')
with io.open(dump_filename) as f:
responses = json.load(f)

assert len(responses['httpbin.org']['80'].keys()) == 1

@mocketize(record_truesocket=False)
def test_wrongpath_truesendall(self):
Entry.register(Entry.GET, 'http://httpbin.org/user.agent', Response(status=404))
Expand Down

0 comments on commit 949c6a5

Please sign in to comment.