Update get_device_instance and test_unsupported following review

sdb9696 · sdb9696 · commit 39644157f510 · 2023-11-30T08:05:29.000Z
diff --git a/kasa/aesprotocol.py b/kasa/aesprotocol.py
@@ -64,7 +64,6 @@ def __init__(
         *,
         credentials: Optional[Credentials] = None,
         timeout: Optional[int] = None,
-        login_version: Optional[int] = 1,
     ) -> None:
         super().__init__(host=host, port=self.DEFAULT_PORT)
 
@@ -91,7 +90,6 @@ def __init__(
         self.terminal_uuid = None
         self.http_client: Optional[httpx.AsyncClient] = None
         self.request_id_generator = SnowflakeId(1, 1)
-        self.login_version = login_version
         self.login_token = None
 
         _LOGGER.debug("Created AES object for %s", self.host)
diff --git a/kasa/discover.py b/kasa/discover.py
@@ -21,7 +21,7 @@
 from kasa.json import dumps as json_dumps
 from kasa.json import loads as json_loads
 from kasa.klapprotocol import TPLinkKlap
-from kasa.protocol import TPLinkSmartHomeProtocol
+from kasa.protocol import TPLinkProtocol, TPLinkSmartHomeProtocol
 from kasa.smartdevice import SmartDevice, SmartDeviceException
 from kasa.smartplug import SmartPlug
 from kasa.tapo.tapoplug import TapoPlug
@@ -380,37 +380,38 @@ def _get_device_instance(
                 f"Unable to read response from device: {ip}: {ex}"
             ) from ex
 
-        if discovery_result.mgt_encrypt_schm.encrypt_type in ("KLAP", "AES"):
-            type_ = discovery_result.device_type
-            device_class = None
-            supported_device_types: dict[str, Type[SmartDevice]] = {
-                "SMART.TAPOPLUG": TapoPlug,
-                "SMART.KASAPLUG": TapoPlug,
-                "IOT.SMARTPLUGSWITCH": SmartPlug,
-            }
-
-            if discovery_result.device_type in supported_device_types:
-                device_class = supported_device_types[discovery_result.device_type]
-
-            if device_class:
-                _LOGGER.debug("[DISCOVERY] %s << %s", ip, info)
-                device = device_class(ip, port=port, credentials=credentials)
-                device.update_from_discover_info(discovery_result.get_dict())
-                if discovery_result.mgt_encrypt_schm.encrypt_type == "KLAP":
-                    device.protocol = TPLinkKlap(ip, credentials=credentials)
-                else:
-                    device.protocol = TPLinkAes(
-                        ip,
-                        credentials=credentials,
-                        login_version=discovery_result.mgt_encrypt_schm.lv,
-                    )
-                return device
-            else:
-                raise UnsupportedDeviceException(
-                    f"Unsupported device {ip} of type {type_}: {info}"
-                )
-        else:
-            raise UnsupportedDeviceException(f"Unsupported device {ip}: {info}")
+        type_ = discovery_result.device_type
+        encrypt_type_ = (
+            f"{type_.split('.')[0]}.{discovery_result.mgt_encrypt_schm.encrypt_type}"
+        )
+        device_class = None
+
+        supported_device_types: dict[str, Type[SmartDevice]] = {
+            "SMART.TAPOPLUG": TapoPlug,
+            "SMART.KASAPLUG": TapoPlug,
+            "IOT.SMARTPLUGSWITCH": SmartPlug,
+        }
+        supported_device_protocols: dict[str, Type[TPLinkProtocol]] = {
+            "IOT.KLAP": TPLinkKlap,
+            "SMART.AES": TPLinkAes,
+        }
+
+        if (device_class := supported_device_types.get(type_)) is None:
+            _LOGGER.warning("Got unsupported device type: %s", type_)
+            raise UnsupportedDeviceException(
+                f"Unsupported device {ip} of type {type_}: {info}"
+            )
+        if (protocol_class := supported_device_protocols.get(encrypt_type_)) is None:
+            _LOGGER.warning("Got unsupported device type: %s", encrypt_type_)
+            raise UnsupportedDeviceException(
+                f"Unsupported encryption scheme {ip} of type {encrypt_type_}: {info}"
+            )
+
+        _LOGGER.debug("[DISCOVERY] %s << %s", ip, info)
+        device = device_class(ip, port=port, credentials=credentials)
+        device.protocol = protocol_class(ip, credentials=credentials)
+        device.update_from_discover_info(discovery_result.get_dict())
+        return device
 
 
 class DiscoveryResult(BaseModel):
diff --git a/kasa/tests/test_discovery.py b/kasa/tests/test_discovery.py
@@ -114,7 +114,7 @@ def mock_discover(self):
     "result": {
         "device_id": "xx",
         "owner": "xx",
-        "device_type": "SMART.TAPOBULB",
+        "device_type": "SMART.TAPOXMASTREE",
         "device_model": "P110(EU)",
         "ip": "127.0.0.1",
         "mac": "48-22xxx",
@@ -150,7 +150,7 @@ def mock_discover(self):
     discovery_data = UNSUPPORTED
     with pytest.raises(
         UnsupportedDeviceException,
-        match=f"Unsupported device {host} of type SMART.TAPOBULB: {re.escape(str(UNSUPPORTED))}",
+        match=f"Unsupported device {host} of type SMART.TAPOXMASTREE: {re.escape(str(UNSUPPORTED))}",
     ):
         await Discover.discover_single(host)
 
