Add Device.supports_module() and stop querying HomeKit for iot devices (#1808)

Some legacy iot firmware stops responding when an unsupported module follows `system` in a request, which made every update time out once the HomeKit query was added.

This stops querying HomeKit in `update()` for IOT devices and adds `Device.supports_module()` API to allow checking if some module is supported by the device.

kasa/device.py

@@ -47,7 +47,6 @@ Devices support different functionality that are exposed via
 
 >>> for module_name in dev.modules:
 ...     print(module_name)
-homekit
 Energy
 schedule
 usage
@@ -264,6 +263,10 @@ class Device(ABC):
     def modules(self) -> ModuleMapping[Module]:
         """Return the device modules."""
 
+    async def supports_module(self, module: ModuleName[Module] | str) -> bool:
+        """Return True if the device supports the module, may query the device."""
+        return module in self.modules
+
     @property
     @abstractmethod
     def is_on(self) -> bool:

kasa/iot/iotdevice.py

@@ -19,13 +19,13 @@ import inspect
 import logging
 from collections.abc import Callable, Mapping, Sequence
 from datetime import datetime, timedelta, tzinfo
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, ClassVar, cast
 from warnings import warn
 
 from ..device import Device, DeviceInfo, WifiNetwork
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
-from ..exceptions import KasaException, UnsupportedDeviceError
+from ..exceptions import KasaException, TimeoutError, UnsupportedDeviceError
 from ..feature import Feature
 from ..module import Module
 from ..modulemapping import ModuleMapping, ModuleName
@@ -182,6 +182,13 @@ class IotDevice(Device):
 
     emeter_type = "emeter"
 
+    #: Modules that are only queried by supports_module(), not by update().
+    _OPTIONAL_MODULES: ClassVar[
+        dict[str | ModuleName[Any], tuple[type[IotModule], str]]
+    ] = {
+        Module.IotHomeKit: (HomeKit, "smartlife.iot.homekit"),
+    }
+
     def __init__(
         self,
         host: str,
@@ -198,6 +205,7 @@ class IotDevice(Device):
         self._children: Mapping[str, IotDevice] = {}
         self._modules: dict[str | ModuleName[Module], IotModule] = {}
         self._on_since: datetime | None = None
+        self._optional_module_support: dict[str | ModuleName[Module], bool] = {}
 
     @property
     def children(self) -> Sequence[IotDevice]:
@@ -330,8 +338,6 @@ class IotDevice(Device):
 
     async def _initialize_modules(self) -> None:
         """Initialize modules not added in init."""
-        self.add_module(Module.IotHomeKit, HomeKit(self, "smartlife.iot.homekit"))
-
         if self.has_emeter:
             _LOGGER.debug(
                 "The device has emeter, querying its information along sysinfo"
@@ -444,6 +450,27 @@ class IotDevice(Device):
 
             self._supported_modules = supported
 
+    async def supports_module(self, module: ModuleName[Module] | str) -> bool:
+        """Return True if the device supports the module.
+
+        Optional modules are checked with a separate query, the result is remembered.
+        """
+        if (optional := self._OPTIONAL_MODULES.get(module)) is None:
+            return await super().supports_module(module)
+
+        if module not in self._optional_module_support:
+            module_class, key = optional
+            query = module_class(self, key).query()
+            try:
+                response = await self.protocol.query(query, retry_count=0)
+            except TimeoutError:
+                return False
+            data = response.get(key)
+            self._optional_module_support[module] = (
+                isinstance(data, dict) and "err_code" not in data
+            )
+        return self._optional_module_support[module]
+
     def update_from_discover_info(self, info: dict[str, Any]) -> None:
         """Update state from info from the discover call."""
         self._discovery_info = info

kasa/iot/iotstrip.py

@@ -29,7 +29,6 @@ from .modules import (
     Cloud,
     Countdown,
     Emeter,
-    HomeKit,
     Led,
     Schedule,
     Time,
@@ -122,7 +121,6 @@ class IotStrip(IotDevice):
         self.add_module(Module.IotCountdown, Countdown(self, "countdown"))
         self.add_module(Module.Led, Led(self, "system"))
         self.add_module(Module.IotCloud, Cloud(self, "cnCloud"))
-        self.add_module(Module.IotHomeKit, HomeKit(self, "smartlife.iot.homekit"))
         if self.has_emeter:
             _LOGGER.debug(
                 "The device has emeter, querying its information along sysinfo"

tests/iot/modules/test_homekit.py

@@ -8,13 +8,23 @@ from kasa.iot.modules.homekit import HomeKit
 
 from ...device_fixtures import device_iot
 
+HOMEKIT = "smartlife.iot.homekit"
+
+
+async def _homekit_module(dev: IotDevice) -> HomeKit:
+    """Create the module with the data from the device, skip if not supported."""
+    if not await dev.supports_module(Module.IotHomeKit):
+        pytest.skip("HomeKit module not present on this device")
+    homekit = HomeKit(dev, HOMEKIT)
+    dev._last_update.update(await dev.protocol.query(homekit.query()))
+    homekit._initialize_features()
+    return homekit
+
 
 @device_iot
-def test_homekit_getters(dev: IotDevice) -> None:
+async def test_homekit_getters(dev: IotDevice) -> None:
     # HomeKit can be present on any IOT device
-    if Module.IotHomeKit not in dev.modules:
+    homekit = await _homekit_module(dev)
-        pytest.skip("HomeKit module not present on this device")
-    homekit: HomeKit = dev.modules[Module.IotHomeKit]
     info = homekit.info
     if not info:
         pytest.skip("No HomeKit data present for this fixture")
@@ -31,10 +41,8 @@ def test_homekit_getters(dev: IotDevice) -> None:
 
 
 @device_iot
-def test_homekit_feature(dev: IotDevice) -> None:
+async def test_homekit_feature(dev: IotDevice) -> None:
-    if Module.IotHomeKit not in dev.modules:
+    homekit = await _homekit_module(dev)
-        pytest.skip("HomeKit module not present on this device")
-    homekit: HomeKit = dev.modules[Module.IotHomeKit]
     if not homekit.info:
         pytest.skip("No HomeKit data present for this device")
     feature = homekit._all_features.get("homekit_setup_code")
@@ -46,9 +54,7 @@ def test_homekit_feature(dev: IotDevice) -> None:
 
 @device_iot
 def test_initialize_features_skips_when_no_data(dev: IotDevice) -> None:
-    if Module.IotHomeKit not in dev.modules:
+    homekit = HomeKit(dev, HOMEKIT)
-        pytest.skip("HomeKit module not present on this device")
-    homekit: HomeKit = dev.modules[Module.IotHomeKit]
     if "homekit_setup_code" in homekit._all_features:
         pytest.skip("HomeKit feature already present on this device")
     # Patch .data so it looks like no homekit data is present
@@ -57,3 +63,4 @@ def test_initialize_features_skips_when_no_data(dev: IotDevice) -> None:
         homekit._initialize_features()
     # Since there was no data, no features should be added
     assert "homekit_setup_code" not in homekit._all_features
+    assert homekit.info == {}

tests/iot/test_iotdevice.py

@@ -18,12 +18,15 @@ from voluptuous import (
 )
 
 from kasa import DeviceType, KasaException, Module
+from kasa.exceptions import AuthenticationError, TimeoutError
 from kasa.iot import IotDevice
 from kasa.iot.iotmodule import _merge_dict
 from tests.conftest import get_device_for_fixture_protocol, handle_turn_on, turn_on
 from tests.device_fixtures import device_iot, has_emeter_iot, no_emeter_iot
 from tests.fakeprotocol_iot import FakeIotProtocol
 
+HOMEKIT = "smartlife.iot.homekit"
+
 TZ_SCHEMA = Schema(
     {"zone_str": str, "dst_offset": int, "index": All(int, Range(min=0)), "tz_str": str}
 )
@@ -310,3 +313,66 @@ def test_merge_dict() -> None:
             "get_daystat": {"month": 8, "year": 2024},
         }
     }
+
+
+@pytest.mark.requires_dummy
+@device_iot
+async def test_update_does_not_query_homekit(dev, mocker: MockerFixture) -> None:
+    spy = mocker.spy(dev.protocol, "query")
+    await dev.update()
+    assert not [call for call in spy.call_args_list if HOMEKIT in call.args[0]]
+
+
+@pytest.mark.requires_dummy
+@pytest.mark.parametrize(
+    ("fixture", "supported"),
+    [("EP25(US)_1.0_1.0.14.json", True), ("HS100(US)_2.0_1.5.6.json", False)],
+)
+async def test_supports_module(
+    mocker: MockerFixture, fixture: str, supported: bool
+) -> None:
+    dev = await get_device_for_fixture_protocol(fixture, "IOT")
+    spy = mocker.spy(dev.protocol, "query")
+
+    assert await dev.supports_module(Module.IotHomeKit) is supported
+
+    # The answer is remembered, so the second call does not query the device
+    assert await dev.supports_module(Module.IotHomeKit) is supported
+    assert spy.call_count == 1
+
+
+@pytest.mark.requires_dummy
+async def test_timeout_is_not_supported_and_not_remembered(
+    mocker: MockerFixture,
+) -> None:
+    dev = await get_device_for_fixture_protocol("EP25(US)_1.0_1.0.14.json", "IOT")
+    query = dev.protocol.query
+
+    mocker.patch.object(dev.protocol, "query", side_effect=TimeoutError("timeout"))
+    assert await dev.supports_module(Module.IotHomeKit) is False
+
+    mocker.patch.object(dev.protocol, "query", side_effect=query)
+    assert await dev.supports_module(Module.IotHomeKit) is True
+
+
+@pytest.mark.requires_dummy
+async def test_unsupported_answer_is_remembered(mocker: MockerFixture) -> None:
+    dev = await get_device_for_fixture_protocol("HS100(US)_2.0_1.5.6.json", "IOT")
+    query = mocker.patch.object(
+        dev.protocol,
+        "query",
+        return_value={HOMEKIT: {"err_code": -1, "err_msg": "module not support"}},
+    )
+
+    assert await dev.supports_module(Module.IotHomeKit) is False
+    assert await dev.supports_module(Module.IotHomeKit) is False
+    assert query.call_count == 1
+
+
+@pytest.mark.requires_dummy
+async def test_other_errors_propagate(mocker: MockerFixture) -> None:
+    dev = await get_device_for_fixture_protocol("EP25(US)_1.0_1.0.14.json", "IOT")
+    mocker.patch.object(dev.protocol, "query", side_effect=AuthenticationError("auth"))
+
+    with pytest.raises(AuthenticationError):
+        await dev.supports_module(Module.IotHomeKit)

tests/test_device.py

@@ -11,6 +11,7 @@ from contextlib import AbstractContextManager, nullcontext
 from unittest.mock import AsyncMock, patch
 
 import pytest
+from pytest_mock import MockerFixture
 
 import kasa
 from kasa import Credentials, Device, DeviceConfig, DeviceType, KasaException, Module
@@ -77,6 +78,16 @@ async def test_alias(dev: Device) -> None:
     assert dev.alias == original
 
 
+async def test_supports_module(dev: Device, mocker: MockerFixture) -> None:
+    """Test modules of the device are supported without querying the device."""
+    spy = mocker.spy(dev.protocol, "query")
+
+    assert await dev.supports_module("not_a_module") is False
+    for name in dev.modules:
+        assert await dev.supports_module(name) is True
+    spy.assert_not_called()
+
+
 @device_classes
 async def test_device_class_ctors(device_class_name_obj) -> None:
     """Make sure constructor api not broken for new and existing SmartDevices."""