Add retries to protocol queries (#65)

* Add retries to query(), defaults to 3 + add tests

* Catch also json decoding errors for retries

* add missing exceptions file, fix old protocol tests

kasa/tests/test_protocol.py

@@ -1,73 +1,95 @@
 import json
-from unittest import TestCase
 
+import pytest
+
+from ..exceptions import SmartDeviceException
 from ..protocol import TPLinkSmartHomeProtocol
 
 
-class TestTPLinkSmartHomeProtocol(TestCase):
-    def test_encrypt(self):
-        d = json.dumps({"foo": 1, "bar": 2})
-        encrypted = TPLinkSmartHomeProtocol.encrypt(d)
-        # encrypt adds a 4 byte header
-        encrypted = encrypted[4:]
-        self.assertEqual(d, TPLinkSmartHomeProtocol.decrypt(encrypted))
+@pytest.mark.parametrize("retry_count", [1, 3, 5])
+async def test_protocol_retries(mocker, retry_count):
+    def aio_mock_writer(_, __):
+        reader = mocker.patch("asyncio.StreamReader")
+        writer = mocker.patch("asyncio.StreamWriter")
 
-    def test_encrypt_unicode(self):
-        d = "{'snowman': '\u2603'}"
-
-        e = bytes(
-            [
-                208,
-                247,
-                132,
-                234,
-                133,
-                242,
-                159,
-                254,
-                144,
-                183,
-                141,
-                173,
-                138,
-                104,
-                240,
-                115,
-                84,
-                41,
-            ]
+        mocker.patch(
+            "asyncio.StreamWriter.write", side_effect=Exception("dummy exception")
         )
 
-        encrypted = TPLinkSmartHomeProtocol.encrypt(d)
-        # encrypt adds a 4 byte header
-        encrypted = encrypted[4:]
+        return reader, writer
 
-        self.assertEqual(e, encrypted)
+    conn = mocker.patch("asyncio.open_connection", side_effect=aio_mock_writer)
+    with pytest.raises(SmartDeviceException):
+        await TPLinkSmartHomeProtocol.query("127.0.0.1", {}, retry_count=retry_count)
 
-    def test_decrypt_unicode(self):
-        e = bytes(
-            [
-                208,
-                247,
-                132,
-                234,
-                133,
-                242,
-                159,
-                254,
-                144,
-                183,
-                141,
-                173,
-                138,
-                104,
-                240,
-                115,
-                84,
-                41,
-            ]
-        )
+    assert conn.call_count == retry_count + 1
 
-        d = "{'snowman': '\u2603'}"
 
-        self.assertEqual(d, TPLinkSmartHomeProtocol.decrypt(e))
+def test_encrypt():
+    d = json.dumps({"foo": 1, "bar": 2})
+    encrypted = TPLinkSmartHomeProtocol.encrypt(d)
+    # encrypt adds a 4 byte header
+    encrypted = encrypted[4:]
+    assert d == TPLinkSmartHomeProtocol.decrypt(encrypted)
+
+
+def test_encrypt_unicode():
+    d = "{'snowman': '\u2603'}"
+
+    e = bytes(
+        [
+            208,
+            247,
+            132,
+            234,
+            133,
+            242,
+            159,
+            254,
+            144,
+            183,
+            141,
+            173,
+            138,
+            104,
+            240,
+            115,
+            84,
+            41,
+        ]
+    )
+
+    encrypted = TPLinkSmartHomeProtocol.encrypt(d)
+    # encrypt adds a 4 byte header
+    encrypted = encrypted[4:]
+
+    assert e == encrypted
+
+
+def test_decrypt_unicode():
+    e = bytes(
+        [
+            208,
+            247,
+            132,
+            234,
+            133,
+            242,
+            159,
+            254,
+            144,
+            183,
+            141,
+            173,
+            138,
+            104,
+            240,
+            115,
+            84,
+            41,
+        ]
+    )
+
+    d = "{'snowman': '\u2603'}"
+
+    assert d == TPLinkSmartHomeProtocol.decrypt(e)