Custom encoding and fix for chunked responses

mocket/compat.py

@@ -26,21 +26,21 @@
 encoding = os.getenv("MOCKET_ENCODING", 'utf-8')
 
 
-def encode_utf8(s):
+def encode_to_bytes(s, charset=encoding):
     if isinstance(s, text_type):
-        s = s.encode(encoding)
+        s = s.encode(charset)
     return byte_type(s)
 
 
-def decode_utf8(s):
+def decode_from_bytes(s, charset=encoding):
     if isinstance(s, byte_type):
-        s = s.decode(encoding)
+        s = s.decode(charset)
     return text_type(s)
 
 
 def shsplit(s):
     if PY2:
-        s = encode_utf8(s)
+        s = encode_to_bytes(s)
     else:
-        s = decode_utf8(s)
+        s = decode_from_bytes(s)
     return shlex.split(s)

mocket/mocket.py

@@ -7,19 +7,19 @@
 import io
 import collections
 import hashlib
-import gzip
+import select
 from datetime import datetime, timedelta
 
 import decorator
+import hexdump
 
 from .compat import (
-    encode_utf8,
-    decode_utf8,
+    encode_to_bytes,
+    decode_from_bytes,
     basestring,
     byte_type,
     text_type,
     FileNotFoundError,
-    encoding,
     JSONDecodeError,
 )
 
@@ -113,7 +113,7 @@ def __init__(self, family=socket.AF_INET, type=socket.SOCK_STREAM, proto=0):
         self.fd = io.BytesIO()
         self._closed = True
         self._connected = False
-        self._buflen = 1024
+        self._buflen = 65536
         self._entry = None
         self.family = family
         self.type = type
@@ -210,9 +210,9 @@ def _connect(self):  # pragma: no cover
             self._connected = True
 
     def true_sendall(self, data, *args, **kwargs):
-        req = decode_utf8(data)
+        req = decode_from_bytes(data)
         # make request unique again
-        req_signature = hashlib.md5(encode_utf8(''.join(sorted(req.split('\r\n'))))).hexdigest()
+        req_signature = hashlib.md5(encode_to_bytes(''.join(sorted(req.split('\r\n'))))).hexdigest()
         # port should be always a string
         port = text_type(self._port)
 
@@ -243,58 +243,33 @@ def true_sendall(self, data, *args, **kwargs):
 
         # try to get the response from the dictionary
         try:
-            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)
-                    try:
-                        gzip_file.write(line)
-                    finally:
-                        gzip_file.close()
-                    line = gzip_buffer.getvalue()
-                r_lines.append(line)
-            encoded_response = b'\r\n'.join(r_lines)
-
+            try:
+                encoded_response = hexdump.restore(response_dict['response'])
+            except TypeError:  # pragma: no cover
+                # Python 2
+                encoded_response = hexdump.restore(encode_to_bytes(response_dict['response']))
         # if not available, call the real sendall
         except KeyError:
             self._connect()
             self.true_socket.sendall(data, *args, **kwargs)
-            r = io.BytesIO()
-            while True:
+            encoded_response = b''
+            # https://github.com/kennethreitz/requests/blob/master/tests/testserver/server.py#L13
+            while select.select([self.true_socket], [], [], 0.5)[0]:
                 recv = self.true_socket.recv(self._buflen)
-                if r.write(recv) < self._buflen:
+                if recv:
+                    encoded_response += recv
+                else:
                     break
 
-            encoded_response = r.getvalue()
-
             # dump the resulting dictionary to a JSON file
             if Mocket.get_truesocket_recording_dir():
-                response_dict['request'] = req
-                lines = response_dict['response'] = []
-                gzipped_lines = response_dict['gzip'] = []
 
-                # update the dictionary with the response obtained
-                for line_no, line in enumerate(encoded_response.split(b'\r\n')):
-
-                    try:
-                        line = decode_utf8(line)
-                    except UnicodeDecodeError:
-                        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)
+                # update the dictionary with request and response lines
+                response_dict['request'] = req
+                response_dict['response'] = hexdump.dump(encoded_response)
 
-                with io.open(path, mode='w', encoding=encoding) as f:
-                    f.write(decode_utf8(json.dumps(responses, indent=4, sort_keys=True)))
+                with io.open(path, mode='w') as f:
+                    f.write(decode_from_bytes(json.dumps(responses, indent=4, sort_keys=True)))
 
         # response back to .sendall() which writes it to the mocket socket and flush the BytesIO
         return encoded_response
@@ -416,12 +391,12 @@ def __init__(self, location, responses):
         for r in responses:
             if not getattr(r, 'data', False):
                 if isinstance(r, text_type):
-                    r = encode_utf8(r)
+                    r = encode_to_bytes(r)
                 r = self.response_cls(r)
             lresponses.append(r)
         else:
             if not responses:
-                lresponses = [self.response_cls(encode_utf8(''))]
+                lresponses = [self.response_cls(encode_to_bytes(''))]
         self.responses = lresponses
 
     def can_handle(self, data):

mocket/mockhttp.py

@@ -6,7 +6,7 @@
 
 import magic
 
-from .compat import BaseHTTPRequestHandler, urlsplit, parse_qs, encode_utf8, decode_utf8
+from .compat import BaseHTTPRequestHandler, urlsplit, parse_qs, encode_to_bytes, decode_from_bytes
 from .mocket import Mocket, MocketEntry
 
 
@@ -16,8 +16,8 @@
 
 class Request(BaseHTTPRequestHandler):
     def __init__(self, data):
-        _, self.body = decode_utf8(data).split('\r\n\r\n', 1)
-        self.rfile = BytesIO(encode_utf8(data))
+        _, self.body = decode_from_bytes(data).split('\r\n\r\n', 1)
+        self.rfile = BytesIO(encode_to_bytes(data))
         self.raw_requestline = self.rfile.readline()
         self.error_code = self.error_message = None
         self.parse_request()
@@ -33,7 +33,7 @@ def __init__(self, body='', status=200, headers=None):
             self.body = body.read()
             is_file_object = True
         except AttributeError:
-            self.body = encode_utf8(body)
+            self.body = encode_to_bytes(body)
         self.status = status
         self.headers = {
             'Status': str(self.status),
@@ -45,7 +45,7 @@ def __init__(self, body='', status=200, headers=None):
         if not is_file_object:
             self.headers['Content-Type'] = 'text/plain; charset=utf-8'
         else:
-            self.headers['Content-Type'] = decode_utf8(magic.from_buffer(self.body, mime=True))
+            self.headers['Content-Type'] = decode_from_bytes(magic.from_buffer(self.body, mime=True))
         for k, v in headers.items():
             self.headers['-'.join([token.capitalize() for token in k.split('-')])] = v
         self.data = self.get_protocol_data() + self.body
@@ -97,7 +97,7 @@ def can_handle(self, data):
         True
         """
         try:
-            requestline, _ = decode_utf8(data).split(CRLF, 1)
+            requestline, _ = decode_from_bytes(data).split(CRLF, 1)
             method, path, version = self._parse_requestline(requestline)
         except ValueError:
             return self == Mocket._last_entry

mocket/mockredis.py

@@ -2,7 +2,7 @@
 from __future__ import unicode_literals
 from itertools import chain
 
-from .compat import text_type, byte_type, encode_utf8, decode_utf8, shsplit
+from .compat import text_type, byte_type, encode_to_bytes, decode_from_bytes, shsplit
 from .mocket import MocketEntry, Mocket
 
 
@@ -19,15 +19,15 @@ def __init__(self, data=None):
 class Redisizer(byte_type):
     @staticmethod
     def tokens(iterable):
-        iterable = [encode_utf8(x) for x in iterable]
+        iterable = [encode_to_bytes(x) for x in iterable]
         return ['*{0}'.format(len(iterable)).encode('utf-8')] + list(chain(*zip(['${0}'.format(len(x)).encode('utf-8') for x in iterable], iterable)))
 
     @staticmethod
     def redisize(data):
         if isinstance(data, Redisizer):
             return data
         if isinstance(data, byte_type):
-            data = decode_utf8(data)
+            data = decode_from_bytes(data)
         CONVERSION = {
             dict: lambda x: b'\r\n'.join(Redisizer.tokens(list(chain(*tuple(x.items()))))),
             int: lambda x: ':{0}'.format(x).encode('utf-8'),

requirements.txt

@@ -1,3 +1,4 @@
 python-magic
 six
-decorator
+decorator
+hexdump

setup.py

@@ -33,7 +33,7 @@
 
 setup(
     name='mocket',
-    version='1.7.4',
+    version='1.7.5',
     # author='Andrea de Marco, Giorgio Salluzzo',
     author='Giorgio Salluzzo',
     # author_email='24erre@gmail.com, giorgio.salluzzo@gmail.com',

tests/main/test_http.py

@@ -68,7 +68,7 @@ def test_truesendall_with_gzip_recording(self):
 
     @mocketize(truesocket_recording_dir=recording_directory)
     def test_truesendall_with_chunk_recording(self):
-        url = 'http://httpbin.org/range/2048?chunk_size=256'
+        url = 'http://httpbin.org/range/70000?chunk_size=65536'
 
         requests.get(url)
         resp = requests.get(url)

tests/main/test_mocket.py

@@ -5,7 +5,7 @@
 import pytest
 
 from mocket import Mocket, mocketize, MocketEntry
-from mocket.compat import encode_utf8
+from mocket.compat import encode_to_bytes
 
 
 class MocketTestCase(TestCase):
@@ -69,11 +69,11 @@ def test_getentry(self):
 
     def test_getresponse(self):
         entry = MocketEntry(('localhost', 8080), ['Show me.\r\n'])
-        self.assertEqual(entry.get_response(), encode_utf8('Show me.\r\n'))
+        self.assertEqual(entry.get_response(), encode_to_bytes('Show me.\r\n'))
 
     def test_empty_getresponse(self):
         entry = MocketEntry(('localhost', 8080), [])
-        self.assertEqual(entry.get_response(), encode_utf8(''))
+        self.assertEqual(entry.get_response(), encode_to_bytes(''))
 
 
 class MocketizeTestCase(TestCase):

tests/main/tests.main.test_http.TrueHttpEntryTestCase.test_truesendall_with_dump_from_recording.json

@@ -1,16 +1,14 @@
 {
-	"httpbin.org": {
-		"80": {
-			"eaa1896a5d6eed6d4dd83a6696efc485": {
-				"gzip": [],
-				"request": "GET /ip HTTP/1.1\r\nHost: httpbin.org\r\nConnection: keep-alive\r\nAccept-Encoding: gzip, deflate\r\nAccept: */*\r\nuser-agent: Fake-User-Agent\r\n\r\n",
-				"response": ["HTTP/1.1 200 OK", "Server: nginx", "Date: Wed, 15 Feb 2017 13:28:10 GMT", "Content-Type: application/json", "Content-Length: 33", "Connection: keep-alive", "Access-Control-Allow-Origin: *", "Access-Control-Allow-Credentials: true", "", "{\n  \"origin\": \"185.27.157.132\"\n}\n"]
-			},
-			"31ae71d6074aa0e73fb4feb91f28546a": {
-				"gzip": [11],
-				"request": "GET /gzip HTTP/1.1\r\nHost: httpbin.org\r\nConnection: keep-alive\r\nAccept-Encoding: gzip, deflate\r\nAccept: */*\r\nuser-agent: Fake-User-Agent\r\n\r\n",
-				"response": ["HTTP/1.1 200 OK", "Server: nginx", "Date: Wed, 15 Feb 2017 13:28:10 GMT", "Content-Type: application/json", "Content-Length: 172", "Connection: keep-alive", "Content-Encoding: gzip", "Access-Control-Allow-Origin: *", "Access-Control-Allow-Credentials: true", "", "{\n  \"gzipped\": true, \n  \"headers\": {\n    \"Accept\": \"*/*\", \n    \"Accept-Encoding\": \"gzip, deflate\", \n    \"Host\": \"httpbin.org\", \n    \"User-Agent\": \"Fake-User-Agent\"\n  }, \n  \"method\": \"GET\", \n  \"origin\": \"185.27.157.132\"\n}\n"]
-			}
-		}
-	}
+    "httpbin.org": {
+        "80": {
+            "31ae71d6074aa0e73fb4feb91f28546a": {
+                "request": "GET /gzip HTTP/1.1\r\nHost: httpbin.org\r\nAccept-Encoding: gzip, deflate\r\nAccept: */*\r\nConnection: keep-alive\r\nuser-agent: Fake-User-Agent\r\n\r\n",
+                "response": "485454502f312e3120323030204f4b0d0a5365727665723a206e67696e780d0a446174653a205765642c2032322046656220323031372031313a33383a313020474d540d0a436f6e74656e742d547970653a206170706c69636174696f6e2f6a736f6e0d0a436f6e74656e742d4c656e6774683a203137300d0a436f6e6e656374696f6e3a206b6565702d616c6976650d0a436f6e74656e742d456e636f64696e673a20677a69700d0a4163636573732d436f6e74726f6c2d416c6c6f772d4f726967696e3a202a0d0a4163636573732d436f6e74726f6c2d416c6c6f772d43726564656e7469616c733a20747275650d0a0d0a1f8b08002278ad5802ff4d8e410ec2201045f79c82ccb22998eaa2d65d17550fa007a865a4a40a047063d3bb0b34695cfef7df9fcc4c2805f955d6a280130dee83254d6cc45ea0f391cd3146d00e03da103314bb02b2b451d6e9c108a565aad3b5920a7cbefa809b78353e8fc710ec43696e9cdcbabb47c75a893a1be77e42f687a2b3ac3fbd318c26bd0997eeb6aec13825954eac39f07dd5f063cdabba06b2901f8d4f4c9bdc000000"
+            },
+            "eaa1896a5d6eed6d4dd83a6696efc485": {
+                "request": "GET /ip HTTP/1.1\r\nHost: httpbin.org\r\nAccept-Encoding: gzip, deflate\r\nAccept: */*\r\nConnection: keep-alive\r\nuser-agent: Fake-User-Agent\r\n\r\n",
+                "response": "485454502f312e3120323030204f4b0d0a5365727665723a206e67696e780d0a446174653a205765642c2032322046656220323031372031313a33383a303920474d540d0a436f6e74656e742d547970653a206170706c69636174696f6e2f6a736f6e0d0a436f6e74656e742d4c656e6774683a2033320d0a436f6e6e656374696f6e3a206b6565702d616c6976650d0a4163636573732d436f6e74726f6c2d416c6c6f772d4f726967696e3a202a0d0a4163636573732d436f6e74726f6c2d416c6c6f772d43726564656e7469616c733a20747275650d0a0d0a7b0a2020226f726967696e223a202239332e3231392e38372e313737220a7d0a"
+            }
+        }
+    }
 }