Make device port configurable (#471)

kasa/tests/test_discovery.py

@@ -52,12 +52,14 @@ async def test_type_unknown():
         Discover._get_device_class(invalid_info)
 
 
-async def test_discover_single(discovery_data: dict, mocker):
+@pytest.mark.parametrize("custom_port", [123, None])
+async def test_discover_single(discovery_data: dict, mocker, custom_port):
     """Make sure that discover_single returns an initialized SmartDevice instance."""
     mocker.patch("kasa.TPLinkSmartHomeProtocol.query", return_value=discovery_data)
-    x = await Discover.discover_single("127.0.0.1")
+    x = await Discover.discover_single("127.0.0.1", port=custom_port)
     assert issubclass(x.__class__, SmartDevice)
     assert x._sys_info is not None
+    assert x.port == custom_port
 
 
 INVALIDS = [

kasa/tests/test_protocol.py

@@ -129,6 +129,36 @@ async def test_protocol_logging(mocker, caplog, log_level):
         assert "success" not in caplog.text
 
 
+@pytest.mark.parametrize("custom_port", [123, None])
+async def test_protocol_custom_port(mocker, custom_port):
+    encrypted = TPLinkSmartHomeProtocol.encrypt('{"great":"success"}')[
+        TPLinkSmartHomeProtocol.BLOCK_SIZE :
+    ]
+
+    async def _mock_read(byte_count):
+        nonlocal encrypted
+        if byte_count == TPLinkSmartHomeProtocol.BLOCK_SIZE:
+            return struct.pack(">I", len(encrypted))
+        if byte_count == len(encrypted):
+            return encrypted
+        raise ValueError(f"No mock for {byte_count}")
+
+    def aio_mock_writer(_, port):
+        reader = mocker.patch("asyncio.StreamReader")
+        writer = mocker.patch("asyncio.StreamWriter")
+        if custom_port is None:
+            assert port == 9999
+        else:
+            assert port == custom_port
+        mocker.patch.object(reader, "readexactly", _mock_read)
+        return reader, writer
+
+    protocol = TPLinkSmartHomeProtocol("127.0.0.1", port=custom_port)
+    mocker.patch("asyncio.open_connection", side_effect=aio_mock_writer)
+    response = await protocol.query({})
+    assert response == {"great": "success"}
+
+
 def test_encrypt():
     d = json.dumps({"foo": 1, "bar": 2})
     encrypted = TPLinkSmartHomeProtocol.encrypt(d)