Add support for pairing devices with hubs (#859)

Co-authored-by: Steven B. <51370195+sdb9696@users.noreply.github.com>

kasa/cli/hub.py

@@ -0,0 +1,96 @@
+"""Hub-specific commands."""
+
+import asyncio
+
+import asyncclick as click
+
+from kasa import DeviceType, Module, SmartDevice
+from kasa.smart import SmartChildDevice
+
+from .common import (
+    echo,
+    error,
+    pass_dev,
+)
+
+
+def pretty_category(cat: str):
+    """Return pretty category for paired devices."""
+    return SmartChildDevice.CHILD_DEVICE_TYPE_MAP.get(cat)
+
+
+@click.group()
+@pass_dev
+async def hub(dev: SmartDevice):
+    """Commands controlling hub child device pairing."""
+    if dev.device_type is not DeviceType.Hub:
+        error(f"{dev} is not a hub.")
+
+    if dev.modules.get(Module.ChildSetup) is None:
+        error(f"{dev} does not have child setup module.")
+
+
+@hub.command(name="list")
+@pass_dev
+async def hub_list(dev: SmartDevice):
+    """List hub paired child devices."""
+    for c in dev.children:
+        echo(f"{c.device_id}: {c}")
+
+
+@hub.command(name="supported")
+@pass_dev
+async def hub_supported(dev: SmartDevice):
+    """List supported hub child device categories."""
+    cs = dev.modules[Module.ChildSetup]
+
+    cats = [cat["category"] for cat in await cs.get_supported_device_categories()]
+    for cat in cats:
+        echo(f"Supports: {cat}")
+
+
+@hub.command(name="pair")
+@click.option("--timeout", default=10)
+@pass_dev
+async def hub_pair(dev: SmartDevice, timeout: int):
+    """Pair all pairable device.
+
+    This will pair any child devices currently in pairing mode.
+    """
+    cs = dev.modules[Module.ChildSetup]
+
+    echo(f"Finding new devices for {timeout} seconds...")
+
+    pair_res = await cs.pair(timeout=timeout)
+    if not pair_res:
+        echo("No devices found.")
+
+    for child in pair_res:
+        echo(
+            f'Paired {child["name"]} ({child["device_model"]}, '
+            f'{pretty_category(child["category"])}) with id {child["device_id"]}'
+        )
+
+
+@hub.command(name="unpair")
+@click.argument("device_id")
+@pass_dev
+async def hub_unpair(dev, device_id: str):
+    """Unpair given device."""
+    cs = dev.modules[Module.ChildSetup]
+
+    # Accessing private here, as the property exposes only values
+    if device_id not in dev._children:
+        error(f"{dev} does not have children with identifier {device_id}")
+
+    res = await cs.unpair(device_id=device_id)
+    # Give the device some time to update its internal state, just in case.
+    await asyncio.sleep(1)
+    await dev.update()
+
+    if device_id not in dev._children:
+        echo(f"Unpaired {device_id}")
+    else:
+        error(f"Failed to unpair {device_id}")
+
+    return res

kasa/cli/main.py

@@ -93,6 +93,7 @@ def _legacy_type_to_class(_type: str) -> Any:
         "hsv": "light",
         "temperature": "light",
         "effect": "light",
+        "hub": "hub",
     },
     result_callback=json_formatter_cb,
 )

kasa/feature.py

@@ -76,6 +76,7 @@ from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from .device import Device
+    from .module import Module
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -142,7 +143,7 @@ class Feature:
     #: Callable coroutine or name of the method that allows changing the value
     attribute_setter: str | Callable[..., Coroutine[Any, Any, Any]] | None = None
     #: Container storing the data, this overrides 'device' for getters
-    container: Any = None
+    container: Device | Module | None = None
     #: Icon suggestion
     icon: str | None = None
     #: Attribute containing the name of the unit getter property.

kasa/module.py

@@ -154,6 +154,7 @@ class Module(ABC):
     )
     ChildLock: Final[ModuleName[smart.ChildLock]] = ModuleName("ChildLock")
     TriggerLogs: Final[ModuleName[smart.TriggerLogs]] = ModuleName("TriggerLogs")
+    ChildSetup: Final[ModuleName[smart.ChildSetup]] = ModuleName("ChildSetup")
 
     HomeKit: Final[ModuleName[smart.HomeKit]] = ModuleName("HomeKit")
     Matter: Final[ModuleName[smart.Matter]] = ModuleName("Matter")

kasa/smart/modules/__init__.py

@@ -8,6 +8,7 @@ from .brightness import Brightness
 from .childdevice import ChildDevice
 from .childlock import ChildLock
 from .childprotection import ChildProtection
+from .childsetup import ChildSetup
 from .clean import Clean
 from .cloud import Cloud
 from .color import Color
@@ -47,6 +48,7 @@ __all__ = [
     "DeviceModule",
     "ChildDevice",
     "ChildLock",
+    "ChildSetup",
     "BatterySensor",
     "HumiditySensor",
     "TemperatureSensor",

kasa/smart/modules/childsetup.py

@@ -0,0 +1,84 @@
+"""Implementation for child device setup.
+
+This module allows pairing and disconnecting child devices.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+from ...feature import Feature
+from ..smartmodule import SmartModule
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class ChildSetup(SmartModule):
+    """Implementation for child device setup."""
+
+    REQUIRED_COMPONENT = "child_quick_setup"
+    QUERY_GETTER_NAME = "get_support_child_device_category"
+
+    def _initialize_features(self) -> None:
+        """Initialize features."""
+        self._add_feature(
+            Feature(
+                self._device,
+                id="pair",
+                name="Pair",
+                container=self,
+                attribute_setter="pair",
+                category=Feature.Category.Config,
+                type=Feature.Type.Action,
+            )
+        )
+
+    async def get_supported_device_categories(self) -> list[dict]:
+        """Get supported device categories."""
+        categories = await self.call("get_support_child_device_category")
+        return categories["get_support_child_device_category"]["device_category_list"]
+
+    async def pair(self, *, timeout: int = 10) -> list[dict]:
+        """Scan for new devices and pair after discovering first new device."""
+        await self.call("begin_scanning_child_device")
+
+        _LOGGER.info("Waiting %s seconds for discovering new devices", timeout)
+        await asyncio.sleep(timeout)
+        detected = await self._get_detected_devices()
+
+        if not detected["child_device_list"]:
+            _LOGGER.info("No devices found.")
+            return []
+
+        _LOGGER.info(
+            "Discovery done, found %s devices: %s",
+            len(detected["child_device_list"]),
+            detected,
+        )
+
+        await self._add_devices(detected)
+
+        return detected["child_device_list"]
+
+    async def unpair(self, device_id: str) -> dict:
+        """Remove device from the hub."""
+        _LOGGER.debug("Going to unpair %s from %s", device_id, self)
+
+        payload = {"child_device_list": [{"device_id": device_id}]}
+        return await self.call("remove_child_device_list", payload)
+
+    async def _add_devices(self, devices: dict) -> dict:
+        """Add devices based on get_detected_device response.
+
+        Pass the output from :ref:_get_detected_devices: as a parameter.
+        """
+        res = await self.call("add_child_device_list", devices)
+        return res
+
+    async def _get_detected_devices(self) -> dict:
+        """Return list of devices detected during scanning."""
+        param = {"scan_list": await self.get_supported_device_categories()}
+        res = await self.call("get_scan_child_device_list", param)
+        _LOGGER.debug("Scan status: %s", res)
+        return res["get_scan_child_device_list"]

kasa/smart/smartdevice.py

@@ -537,6 +537,21 @@ class SmartDevice(Device):
             )
         )
 
+        if self.parent is not None and (
+            cs := self.parent.modules.get(Module.ChildSetup)
+        ):
+            self._add_feature(
+                Feature(
+                    device=self,
+                    id="unpair",
+                    name="Unpair device",
+                    container=cs,
+                    attribute_setter=lambda: cs.unpair(self.device_id),
+                    category=Feature.Category.Debug,
+                    type=Feature.Type.Action,
+                )
+            )
+
         for module in self.modules.values():
             module._initialize_features()
             for feat in module._module_features.values():

tests/cli/test_hub.py

@@ -0,0 +1,53 @@
+import pytest
+from pytest_mock import MockerFixture
+
+from kasa import DeviceType, Module
+from kasa.cli.hub import hub
+
+from ..device_fixtures import HUBS_SMART, hubs_smart, parametrize, plug_iot
+
+
+@hubs_smart
+async def test_hub_pair(dev, mocker: MockerFixture, runner, caplog):
+    """Test that pair calls the expected methods."""
+    cs = dev.modules.get(Module.ChildSetup)
+    # Patch if the device supports the module
+    if cs is not None:
+        mock_pair = mocker.patch.object(cs, "pair")
+
+    res = await runner.invoke(hub, ["pair"], obj=dev, catch_exceptions=False)
+    if cs is None:
+        assert "is not a hub" in res.output
+        return
+
+    mock_pair.assert_awaited()
+    assert "Finding new devices for 10 seconds" in res.output
+    assert res.exit_code == 0
+
+
+@parametrize("hubs smart", model_filter=HUBS_SMART, protocol_filter={"SMART"})
+async def test_hub_unpair(dev, mocker: MockerFixture, runner):
+    """Test that unpair calls the expected method."""
+    if not dev.children:
+        pytest.skip("Cannot test without child devices")
+
+    id_ = next(iter(dev.children)).device_id
+
+    cs = dev.modules.get(Module.ChildSetup)
+    mock_unpair = mocker.spy(cs, "unpair")
+
+    res = await runner.invoke(hub, ["unpair", id_], obj=dev, catch_exceptions=False)
+
+    mock_unpair.assert_awaited()
+    assert f"Unpaired {id_}" in res.output
+    assert res.exit_code == 0
+
+
+@plug_iot
+async def test_non_hub(dev, mocker: MockerFixture, runner):
+    """Test that hub commands return an error if executed on a non-hub."""
+    assert dev.device_type is not DeviceType.Hub
+    res = await runner.invoke(
+        hub, ["unpair", "dummy_id"], obj=dev, catch_exceptions=False
+    )
+    assert "is not a hub" in res.output

tests/conftest.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import os
 import sys
 import warnings
 from pathlib import Path
@@ -8,6 +9,9 @@ from unittest.mock import MagicMock, patch
 
 import pytest
 
+# TODO: this and runner fixture could be moved to tests/cli/conftest.py
+from asyncclick.testing import CliRunner
+
 from kasa import (
     DeviceConfig,
     SmartProtocol,
@@ -149,3 +153,12 @@ def mock_datagram_endpoint(request):  # noqa: PT004
             side_effect=_create_datagram_endpoint,
         ):
             yield
+
+
+@pytest.fixture
+def runner():
+    """Runner fixture that unsets the KASA_ environment variables for tests."""
+    KASA_VARS = {k: None for k, v in os.environ.items() if k.startswith("KASA_")}
+    runner = CliRunner(env=KASA_VARS)
+
+    return runner

tests/fakeprotocol_smart.py

@@ -171,6 +171,16 @@ class FakeSmartTransport(BaseTransport):
                 "setup_payload": "00:0000000-0000.00.000",
             },
         ),
+        # child setup
+        "get_support_child_device_category": (
+            "child_quick_setup",
+            {"device_category_list": [{"category": "subg.trv"}]},
+        ),
+        # no devices found
+        "get_scan_child_device_list": (
+            "child_quick_setup",
+            {"child_device_list": [{"dummy": "response"}], "scan_status": "idle"},
+        ),
     }
 
     def _missing_result(self, method):
@@ -548,6 +558,17 @@ class FakeSmartTransport(BaseTransport):
 
         return {"error_code": 0}
 
+    def _hub_remove_device(self, info, params):
+        """Remove hub device."""
+        items_to_remove = [dev["device_id"] for dev in params["child_device_list"]]
+        children = info["get_child_device_list"]["child_device_list"]
+        new_children = [
+            dev for dev in children if dev["device_id"] not in items_to_remove
+        ]
+        info["get_child_device_list"]["child_device_list"] = new_children
+
+        return {"error_code": 0}
+
     def get_child_device_queries(self, method, params):
         return self._get_method_from_info(method, params)
 
@@ -658,8 +679,15 @@ class FakeSmartTransport(BaseTransport):
             return self._set_on_off_gradually_info(info, params)
         elif method == "set_child_protection":
             return self._update_sysinfo_key(info, "child_protection", params["enable"])
-        # Vacuum special actions
-        elif method in ["playSelectAudio"]:
+        elif method == "remove_child_device_list":
+            return self._hub_remove_device(info, params)
+        # actions
+        elif method in [
+            "begin_scanning_child_device",  # hub pairing
+            "add_child_device_list",  # hub pairing
+            "remove_child_device_list",  # hub pairing
+            "playSelectAudio",  # vacuum special actions
+        ]:
             return {"error_code": 0}
         elif method[:3] == "set":
             target_method = f"get{method[3:]}"

tests/fixtures/smart/H100(EU)_1.0_1.5.10.json

@@ -472,6 +472,24 @@
         "setup_code": "00000000000",
         "setup_payload": "00:0000000000000000000"
     },
+    "get_scan_child_device_list": {
+        "child_device_list": [
+            {
+                "category": "subg.trigger.temp-hmdt-sensor",
+                "device_id": "REDACTED_1",
+                "device_model": "T315",
+                "name": "REDACTED_1"
+            },
+            {
+                "category": "subg.trigger.contact-sensor",
+                "device_id": "REDACTED_2",
+                "device_model": "T110",
+                "name": "REDACTED_2"
+            }
+        ],
+        "scan_status": "scanning",
+        "scan_wait_time": 28
+    },
     "get_support_alarm_type_list": {
         "alarm_type_list": [
             "Doorbell Ring 1",

tests/smart/modules/test_childsetup.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import logging
+
+import pytest
+from pytest_mock import MockerFixture
+
+from kasa import Feature, Module, SmartDevice
+
+from ...device_fixtures import parametrize
+
+childsetup = parametrize(
+    "supports pairing", component_filter="child_quick_setup", protocol_filter={"SMART"}
+)
+
+
+@childsetup
+async def test_childsetup_features(dev: SmartDevice):
+    """Test the exposed features."""
+    cs = dev.modules.get(Module.ChildSetup)
+    assert cs
+
+    assert "pair" in cs._module_features
+    pair = cs._module_features["pair"]
+    assert pair.type == Feature.Type.Action
+
+
+@childsetup
+async def test_childsetup_pair(
+    dev: SmartDevice, mocker: MockerFixture, caplog: pytest.LogCaptureFixture
+):
+    """Test device pairing."""
+    caplog.set_level(logging.INFO)
+    mock_query_helper = mocker.spy(dev, "_query_helper")
+    mocker.patch("asyncio.sleep")
+
+    cs = dev.modules.get(Module.ChildSetup)
+    assert cs
+
+    await cs.pair()
+
+    mock_query_helper.assert_has_awaits(
+        [
+            mocker.call("begin_scanning_child_device", None),
+            mocker.call("get_support_child_device_category", None),
+            mocker.call("get_scan_child_device_list", params=mocker.ANY),
+            mocker.call("add_child_device_list", params=mocker.ANY),
+        ]
+    )
+    assert "Discovery done" in caplog.text
+
+
+@childsetup
+async def test_childsetup_unpair(
+    dev: SmartDevice, mocker: MockerFixture, caplog: pytest.LogCaptureFixture
+):
+    """Test unpair."""
+    mock_query_helper = mocker.spy(dev, "_query_helper")
+    DUMMY_ID = "dummy_id"
+
+    cs = dev.modules.get(Module.ChildSetup)
+    assert cs
+
+    await cs.unpair(DUMMY_ID)
+
+    mock_query_helper.assert_awaited_with(
+        "remove_child_device_list",
+        params={"child_device_list": [{"device_id": DUMMY_ID}]},
+    )

tests/smart/test_smartdevice.py

@@ -988,3 +988,24 @@ async def test_dynamic_devices(dev: Device, caplog: pytest.LogCaptureFixture):
         await dev.update()
 
     assert "Could not find child id for device" not in caplog.text
+
+
+@hubs_smart
+async def test_unpair(dev: SmartDevice, mocker: MockerFixture):
+    """Verify that unpair calls childsetup module."""
+    if not dev.children:
+        pytest.skip("device has no children")
+
+    child = dev.children[0]
+
+    assert child.parent is not None
+    assert Module.ChildSetup in dev.modules
+    cs = dev.modules[Module.ChildSetup]
+
+    unpair_call = mocker.spy(cs, "unpair")
+
+    unpair_feat = child.features.get("unpair")
+    assert unpair_feat
+    await unpair_feat.set_value(None)
+
+    unpair_call.assert_called_with(child.device_id)

tests/test_cli.py

@@ -1,5 +1,4 @@
 import json
-import os
 import re
 from datetime import datetime
 from unittest.mock import ANY, PropertyMock, patch
@@ -62,15 +61,6 @@ from .conftest import (
 pytestmark = [pytest.mark.requires_dummy]
 
 
-@pytest.fixture
-def runner():
-    """Runner fixture that unsets the KASA_ environment variables for tests."""
-    KASA_VARS = {k: None for k, v in os.environ.items() if k.startswith("KASA_")}
-    runner = CliRunner(env=KASA_VARS)
-
-    return runner
-
-
 async def test_help(runner):
     """Test that all the lazy modules are correctly names."""
     res = await runner.invoke(cli, ["--help"])

tests/test_feature.py

@@ -74,7 +74,7 @@ def test_feature_value_container(mocker, dummy_feature: Feature):
         def test_prop(self):
             return "dummy"
 
-    dummy_feature.container = DummyContainer()
+    dummy_feature.container = DummyContainer()  # type: ignore[assignment]
     dummy_feature.attribute_getter = "test_prop"
 
     mock_dev_prop = mocker.patch.object(
@@ -191,7 +191,12 @@ async def test_feature_setters(dev: Device, mocker: MockerFixture):
         exceptions = []
         for feat in dev.features.values():
             try:
-                with patch.object(feat.device.protocol, "query") as query:
+                prot = (
+                    feat.container._device.protocol
+                    if feat.container
+                    else feat.device.protocol
+                )
+                with patch.object(prot, "query", name=feat.id) as query:
                     await _test_feature(feat, query)
             # we allow our own exceptions to avoid mocking valid responses
             except KasaException: