From 8d340f7f199266a26eede5b7265ab6d435e2da4f Mon Sep 17 00:00:00 2001
From: AN Long <aisk1988@gmail.com>
Date: Tue, 14 Nov 2023 18:26:40 +0800
Subject: [PATCH] add tests for aio protocols

---
 tests/test_aio_protocol_binary.py  | 33 ++++++++++++++++++++++++++++++
 tests/test_aio_protocol_compact.py | 33 ++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+)
 create mode 100644 tests/test_aio_protocol_binary.py
 create mode 100644 tests/test_aio_protocol_compact.py

diff --git a/tests/test_aio_protocol_binary.py b/tests/test_aio_protocol_binary.py
new file mode 100644
index 0000000..ac5b4e3
--- /dev/null
+++ b/tests/test_aio_protocol_binary.py
@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+
+from io import BytesIO
+
+import pytest
+
+from thriftpy2.thrift import TType, TPayload
+from thriftpy2.contrib.aio.protocol import binary as proto
+
+
+class TItem(TPayload):
+    thrift_spec = {
+        1: (TType.I32, "id", False),
+        2: (TType.LIST, "phones", (TType.STRING), False),
+    }
+    default_spec = [("id", None), ("phones", None)]
+
+
+class AsyncBytesIO:
+    def __init__(self, b):
+        self.b = b
+
+    async def read(self, *args, **kwargs):
+        return self.b.read(*args, **kwargs)
+
+
+@pytest.mark.asyncio
+async def test_strict_decode():
+    bs = AsyncBytesIO(BytesIO(b"\x00\x00\x00\x0c\x00"  # there is a redundant '\x00'
+                      b"\xe4\xbd\xa0\xe5\xa5\xbd\xe4\xb8\x96\xe7\x95\x8c"))
+    with pytest.raises(UnicodeDecodeError):
+        await proto.read_val(bs, TType.STRING, decode_response=True,
+                             strict_decode=True)
diff --git a/tests/test_aio_protocol_compact.py b/tests/test_aio_protocol_compact.py
new file mode 100644
index 0000000..7ea306c
--- /dev/null
+++ b/tests/test_aio_protocol_compact.py
@@ -0,0 +1,33 @@
+# -*- coding: utf-8 -*-
+
+from io import BytesIO
+
+import pytest
+
+from thriftpy2.thrift import TType, TPayload
+from thriftpy2.contrib.aio.protocol import compact
+from test_aio_protocol_binary import AsyncBytesIO
+
+
+class TItem(TPayload):
+    thrift_spec = {
+        1: (TType.I32, "id", False),
+        2: (TType.LIST, "phones", (TType.STRING), False),
+    }
+    default_spec = [("id", None), ("phones", None)]
+
+
+def gen_proto(bytearray=b''):
+    b = AsyncBytesIO(BytesIO(bytearray))
+    proto = compact.TAsyncCompactProtocol(b)
+    return (b, proto)
+
+
+@pytest.mark.asyncio
+async def test_strict_decode():
+    b, proto = gen_proto(b'\x0c\xe4\xbd\xa0\xe5\xa5\x00'
+                         b'\xbd\xe4\xb8\x96\xe7\x95\x8c')
+    proto.strict_decode = True
+
+    with pytest.raises(UnicodeDecodeError):
+        await proto._read_val(TType.STRING)