test: test bytearray input for SOCKS{4,5}REPLY.loads

sethmlarson · Sep 12, 2022 · 9f992b3 · 9f992b3
1 parent caa0223
commit 9f992b3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 13 deletions.
diff --git a/tests/test_socks4.py b/tests/test_socks4.py
@@ -1,3 +1,4 @@
+from typing import Type, Union
 import pytest
 
 from socksio import (
@@ -99,19 +100,20 @@ def test_socks4_connection_request(command: SOCKS4Command) -> None:
 
 
 @pytest.mark.parametrize("request_reply_code", list(SOCKS4ReplyCode))
-def test_socks4_receive_data(request_reply_code: bytes) -> None:
+@pytest.mark.parametrize("data_type", [bytes, bytearray])
+def test_socks4_receive_data(
+    request_reply_code: bytes, data_type: Union[Type[bytearray], Type[bytes]]
+) -> None:
     conn = SOCKS4Connection(user_id=b"socks")
-
-    reply = conn.receive_data(
-        b"".join(
-            [
-                b"\x00",
-                request_reply_code,
-                (8080).to_bytes(2, byteorder="big"),
-                b"\x7f\x00\x00\x01",
-            ]
-        )
+    data = b"".join(
+        [
+            b"\x00",
+            request_reply_code,
+            (8080).to_bytes(2, byteorder="big"),
+            b"\x7f\x00\x00\x01",
+        ]
     )
+    reply = conn.receive_data(data_type(data))
 
     assert reply == SOCKS4Reply(
         reply_code=SOCKS4ReplyCode(request_reply_code), port=8080, addr="127.0.0.1"

diff --git a/tests/test_socks5.py b/tests/test_socks5.py
@@ -1,3 +1,4 @@
+from typing import Type, Union
 import pytest
 
 from socksio import (
@@ -285,12 +286,14 @@ def test_socks5_request_ipv6(
         ),
     ],
 )
+@pytest.mark.parametrize("data_type", [bytes, bytearray])
 def test_socks5_reply_success(
     authenticated_conn: SOCKS5Connection,
     atype: bytes,
     addr: bytes,
     expected_atype: SOCKS5AType,
     expected_addr: str,
+    data_type: Union[Type[bytes], Type[bytearray]],
 ) -> None:
     data = b"".join(
         [
@@ -302,7 +305,7 @@ def test_socks5_reply_success(
             (1080).to_bytes(2, byteorder="big"),  # port
         ]
     )
-    reply = authenticated_conn.receive_data(data)
+    reply = authenticated_conn.receive_data(data_type(data))
 
     assert authenticated_conn.state == SOCKS5State.TUNNEL_READY
     assert reply == SOCKS5Reply(
@@ -342,13 +345,15 @@ def test_socks5_receive_malformed_data(
         ),
     ],
 )
+@pytest.mark.parametrize("data_type", [bytes, bytearray])
 def test_socks5_reply_error(
     error_code: SOCKS5ReplyCode,
     authenticated_conn: SOCKS5Connection,
     atype: bytes,
     addr: bytes,
     expected_atype: SOCKS5AType,
     expected_addr: str,
+    data_type: Union[Type[bytes], Type[bytearray]],
 ) -> None:
     data = b"".join(
         [
@@ -360,7 +365,7 @@ def test_socks5_reply_error(
             (1080).to_bytes(2, byteorder="big"),  # port
         ]
     )
-    reply = authenticated_conn.receive_data(data)
+    reply = authenticated_conn.receive_data(data_type(data))
 
     assert authenticated_conn.state == SOCKS5State.MUST_CLOSE
     assert reply == SOCKS5Reply(