Add resp encoder class to tests

Author: kristjanvalur

tests/resp.py

@@ -0,0 +1,93 @@
+import itertools
+from types import NoneType
+from typing import Any, Optional
+
+
+class RespEncoder:
+    """
+    A class for simple RESP protocol encodign for unit tests
+    """
+
+    def __init__(self, protocol: int = 2, encoding: str = "utf-8") -> None:
+        self.protocol = protocol
+        self.encoding = encoding
+
+    def encode(self, data: Any, hint: Optional[str] = None) -> bytes:
+        if isinstance(data, dict):
+            if self.protocol > 2:
+                result = f"%{len(data)}\r\n".encode()
+                for key, val in data.items():
+                    result += self.encode(key) + self.encode(val)
+                return result
+            else:
+                # Automatically encode dicts as flattened key, value arrays
+                mylist = list(
+                    itertools.chain(*((key, val) for (key, val) in data.items()))
+                )
+                return self.encode(mylist)
+
+        elif isinstance(data, list):
+            result = f"*{len(data)}\r\n".encode()
+            for val in data:
+                result += self.encode(val)
+            return result
+
+        elif isinstance(data, set):
+            if self.protocol > 2:
+                result = f"~{len(data)}\r\n".encode()
+                for val in data:
+                    result += self.encode(val)
+                return result
+            else:
+                return self.encode(list(data))
+
+        elif isinstance(data, str):
+            enc = data.encode(self.encoding)
+            # long strings or strings with control characters must be encoded as bulk
+            # strings
+            if hint or len(enc) > 20 or b"\r" in enc or b"\n" in enc:
+                return self.encode_bulkstr(enc, hint)
+            return b"+" + enc + b"\r\n"
+
+        elif isinstance(data, bytes):
+            return self.encode_bulkstr(data, hint)
+
+        elif isinstance(data, bool):
+            if self.protocol == 2:
+                return b":1\r\n" if data else b":0\r\n"
+            else:
+                return b"t\r\n" if data else b"f\r\n"
+
+        elif isinstance(data, int):
+            if (data > 2**63 - 1) or (data < -(2**63)):
+                if self.protocol > 2:
+                    return f"({data}\r\n".encode()  # resp3 big int
+                else:
+                    return f"+{data}\r\n".encode()  # force to simple string
+            return f":{data}\r\n".encode()
+        elif isinstance(data, float):
+            if self.protocol > 2:
+                return f",{data}\r\n".encode()  # resp3 double
+            else:
+                return f"+{data}\r\n".encode()  # simple string
+
+        elif isinstance(data, NoneType):
+            if self.protocol > 2:
+                return b"_\r\n"  # resp3 null
+            else:
+                return b"$-1\r\n"  # Null bulk string
+                # some commands return null array: b"*-1\r\n"
+
+        else:
+            raise NotImplementedError
+
+    def encode_bulkstr(self, bstr: bytes, hint: Optional[str]) -> bytes:
+        if self.protocol > 2 and hint is not None:
+            # a resp3 verbatim string
+            return f"={len(bstr)}\r\n{hint}:".encode() + bstr + b"\r\n"
+        else:
+            return f"${len(bstr)}\r\n".encode() + bstr + b"\r\n"
+
+
+def encode(value: Any, protocol: int = 2, hint: Optional[str] = None) -> bytes:
+    return RespEncoder(protocol).encode(value, hint)

tests/test_resp.py

@@ -0,0 +1,97 @@
+from .resp import encode
+
+import pytest
+
+
+@pytest.fixture(params=[2, 3])
+def resp_version(request):
+    return request.param
+
+
+class TestEncoder:
+    def test_simple_str(self):
+        assert encode("foo") == b"+foo\r\n"
+
+    def test_long_str(self):
+        text = "fooling around with the sword in the mud"
+        assert len(text) == 40
+        assert encode(text) == b"$40\r\n" + text.encode() + b"\r\n"
+
+    # test strings with control characters
+    def test_str_with_ctrl_chars(self):
+        text = "foo\r\nbar"
+        assert encode(text) == b"$8\r\nfoo\r\nbar\r\n"
+
+    def test_bytes(self):
+        assert encode(b"foo") == b"$3\r\nfoo\r\n"
+
+    def test_int(self):
+        assert encode(123) == b":123\r\n"
+
+    def test_float(self, resp_version):
+        data = encode(1.23, protocol=resp_version)
+        if resp_version == 2:
+            assert data == b"+1.23\r\n"
+        else:
+            assert data == b",1.23\r\n"
+
+    def test_large_int(self, resp_version):
+        data = encode(2**63, protocol=resp_version)
+        if resp_version == 2:
+            assert data == b"+9223372036854775808\r\n"
+        else:
+            assert data == b"(9223372036854775808\r\n"
+
+    def test_array(self):
+        assert encode([1, 2, 3]) == b"*3\r\n:1\r\n:2\r\n:3\r\n"
+
+    def test_set(self, resp_version):
+        data = encode({1, 2, 3}, protocol=resp_version)
+        if resp_version == 2:
+            assert data == b"*3\r\n:1\r\n:2\r\n:3\r\n"
+        else:
+            assert data == b"~3\r\n:1\r\n:2\r\n:3\r\n"
+
+    def test_map(self, resp_version):
+        data = encode({1: 2, 3: 4}, protocol=resp_version)
+        if resp_version == 2:
+            assert data == b"*4\r\n:1\r\n:2\r\n:3\r\n:4\r\n"
+        else:
+            assert data == b"%2\r\n:1\r\n:2\r\n:3\r\n:4\r\n"
+
+    def test_nested_array(self):
+        assert encode([1, [2, 3]]) == b"*2\r\n:1\r\n*2\r\n:2\r\n:3\r\n"
+
+    def test_nested_map(self, resp_version):
+        data = encode({1: {2: 3}}, protocol=resp_version)
+        if resp_version == 2:
+            assert data == b"*2\r\n:1\r\n*2\r\n:2\r\n:3\r\n"
+        else:
+            assert data == b"%1\r\n:1\r\n%1\r\n:2\r\n:3\r\n"
+
+    def test_null(self, resp_version):
+        data = encode(None, protocol=resp_version)
+        if resp_version == 2:
+            assert data == b"$-1\r\n"
+        else:
+            assert data == b"_\r\n"
+
+    def test_mixed_array(self, resp_version):
+        data = encode([1, "foo", 2.3, None, True], protocol=resp_version)
+        if resp_version == 2:
+            assert data == b"*5\r\n:1\r\n+foo\r\n+2.3\r\n$-1\r\n:1\r\n"
+        else:
+            assert data == b"*5\r\n:1\r\n+foo\r\n,2.3\r\n_\r\nt\r\n"
+
+    def test_bool(self, resp_version):
+        data = encode(True, protocol=resp_version)
+        if resp_version == 2:
+            assert data == b":1\r\n"
+        else:
+            assert data == b"t\r\n"
+
+        data = encode(False, resp_version)
+        if resp_version == 2:
+            assert data == b":0\r\n"
+        else:
+            assert data == b"f\r\n"