fix(core): now encryption schema can be nullable

petretiandrea · Apr 6, 2024 · f7ead0d · f7ead0d
1 parent 181d6fe
commit f7ead0d
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 18 deletions.
diff --git a/plugp100/discovery/__init__.py b/plugp100/discovery/__init__.py
@@ -1,4 +1,4 @@
-from .discovered_device import DiscoveredDevice, EncryptionScheme
+from .discovered_device import DiscoveredDevice, EncryptionSchema
 from .tapo_discovery import TapoDiscovery
 
-__all__ = ["TapoDiscovery", "DiscoveredDevice", "EncryptionScheme"]
+__all__ = ["TapoDiscovery", "DiscoveredDevice", "EncryptionSchema"]
diff --git a/plugp100/discovery/discovered_device.py b/plugp100/discovery/discovered_device.py
@@ -1,4 +1,5 @@
 import dataclasses
+import logging
 from typing import Optional, Any
 
 import aiohttp
@@ -14,7 +15,7 @@ class DiscoveredDevice:
     device_model: str
     ip: str
     mac: str
-    mgt_encrypt_schm: "EncryptionScheme"
+    mgt_encrypt_schm: Optional["EncryptionSchema"]
 
     device_id: Optional[str] = None
     owner: Optional[str] = None
@@ -25,6 +26,10 @@ class DiscoveredDevice:
 
     @staticmethod
     def from_dict(values: dict[str, Any]) -> "DiscoveredDevice":
+        if enc_schema_info := values.get("mgt_encrypt_schm", None):
+            encryption_schema = EncryptionSchema(**enc_schema_info)
+        else:
+            encryption_schema = None
         return DiscoveredDevice(
             device_type=values.get("device_type", values.get("device_type_text")),
             device_model=values.get("device_model", values.get("model")),
@@ -36,7 +41,7 @@ def from_dict(values: dict[str, Any]) -> "DiscoveredDevice":
             is_support_iot_cloud=values.get("is_support_iot_cloud", None),
             obd_src=values.get("obd_src", None),
             factory_default=values.get("factory_default", None),
-            mgt_encrypt_schm=EncryptionScheme(**values.get("mgt_encrypt_schm")),
+            mgt_encrypt_schm=encryption_schema,
         )
 
     @property
@@ -56,25 +61,37 @@ def as_dict(self) -> dict[str, Any]:
                 "encrypt_type": self.mgt_encrypt_schm.encrypt_type,
                 "http_port": self.mgt_encrypt_schm.http_port,
                 "lv": self.mgt_encrypt_schm.lv,
-            },
+            }
+            if self.mgt_encrypt_schm is not None
+            else None,
         }
 
     async def get_tapo_device(
         self, credentials: AuthCredential, session: Optional[aiohttp.ClientSession] = None
     ) -> TapoDevice:
-        config = DeviceConnectConfiguration(
-            host=self.ip,
-            port=self.mgt_encrypt_schm.http_port,
-            credentials=credentials,
-            device_type=self.device_type,
-            encryption_type=self.mgt_encrypt_schm.encrypt_type,
-            encryption_version=self.mgt_encrypt_schm.lv,
-        )
+        if encrypt_schema := self.mgt_encrypt_schm:
+            config = DeviceConnectConfiguration(
+                host=self.ip,
+                port=encrypt_schema.http_port,
+                credentials=credentials,
+                device_type=self.device_type,
+                encryption_type=encrypt_schema.encrypt_type,
+                encryption_version=encrypt_schema.lv,
+            )
+        else:
+            logging.warning(
+                "No encryption schema found for discovered device {}, {}",
+                self.ip,
+                self.device_type,
+            )
+            config = DeviceConnectConfiguration(
+                host=self.ip, port=80, device_type=self.device_type
+            )
         return await connect(config, session)
 
 
 @dataclasses.dataclass
-class EncryptionScheme:
+class EncryptionSchema:
     """Base model for encryption scheme of discovery result."""
 
     is_support_https: Optional[bool] = None

diff --git a/tests/integration/test_power_strip.py b/tests/integration/test_power_strip.py
@@ -1,7 +1,7 @@
 import unittest
 
 from plugp100.new.device_factory import connect
-from plugp100.new.tapoplugstrip import TapoPlugStrip
+from plugp100.new.tapoplug import TapoPlug
 from tests.integration.tapo_test_helper import (
     _test_expose_device_info,
     get_test_config,
@@ -14,7 +14,7 @@ class PowerStripTest(unittest.IsolatedAsyncioTestCase):
 
     async def asyncSetUp(self) -> None:
         connect_config = await get_test_config(device_type="power_strip")
-        self._device: TapoPlugStrip = await connect(connect_config)
+        self._device: TapoPlug = await connect(connect_config)
         await self._device.update()
 
     async def asyncTearDown(self):

diff --git a/tests/integration/test_tapo_discovery.py b/tests/integration/test_tapo_discovery.py
@@ -6,7 +6,7 @@
 from cryptography.hazmat.primitives import serialization, hashes
 from cryptography.hazmat.primitives.asymmetric import padding
 
-from plugp100.discovery.discovered_device import DiscoveredDevice, EncryptionScheme
+from plugp100.discovery.discovered_device import DiscoveredDevice, EncryptionSchema
 from plugp100.discovery.tapo_discovery import (
     TapoDiscovery,
     _build_packet_for_payload_json,
@@ -88,7 +88,7 @@ def test_end_to_end(self):
                     device_model="P105(IT)",
                     ip="192.168.1.3",
                     mac="9C-53-22-A7-C8-35",
-                    mgt_encrypt_schm=EncryptionScheme(
+                    mgt_encrypt_schm=EncryptionSchema(
                         is_support_https=False, encrypt_type="KLAP", http_port=80, lv=2
                     ),
                     device_id="26d9887af76f0d5facf8febeb20f6ec9",