From 9e74e1bd40f53c30f5ff0af3d41d79d6ffabc89c Mon Sep 17 00:00:00 2001
From: Steven B <51370195+sdb9696@users.noreply.github.com>
Date: Mon, 10 Jun 2024 06:21:21 +0100
Subject: [PATCH] Do not add parent only modules to strip sockets (#963)

Excludes modules that child devices report as supported but do not make sense
on a child device like firmware, cloud, time etc.
---
 kasa/smart/smartdevice.py      | 10 ++++++----
 kasa/tests/device_fixtures.py  | 18 ++++++++++++++++++
 kasa/tests/test_childdevice.py | 20 +++++++++++++++++++-
 3 files changed, 43 insertions(+), 5 deletions(-)

diff --git a/kasa/smart/smartdevice.py b/kasa/smart/smartdevice.py
index 0c56dba8..9013fc93 100644
--- a/kasa/smart/smartdevice.py
+++ b/kasa/smart/smartdevice.py
@@ -29,11 +29,11 @@ from .smartmodule import SmartModule
 _LOGGER = logging.getLogger(__name__)
 
 
-# List of modules that wall switches with children, i.e. ks240 report on
+# List of modules that non hub devices with children, i.e. ks240/P300, report on
 # the child but only work on the parent.  See longer note below in _initialize_modules.
 # This list should be updated when creating new modules that could have the
 # same issue, homekit perhaps?
-WALL_SWITCH_PARENT_ONLY_MODULES = [DeviceModule, Time, Firmware, Cloud]
+NON_HUB_PARENT_ONLY_MODULES = [DeviceModule, Time, Firmware, Cloud]
 
 
 # Device must go last as the other interfaces also inherit Device
@@ -196,9 +196,11 @@ class SmartDevice(Device):
         # when they need to be accessed through the children.
         # The logic below ensures that such devices add all but whitelisted, only on
         # the child device.
+        # It also ensures that devices like power strips do not add modules such as
+        # firmware to the child devices.
         skip_parent_only_modules = False
         child_modules_to_skip = {}
-        if self._parent and self._parent.device_type == DeviceType.WallSwitch:
+        if self._parent and self._parent.device_type != DeviceType.Hub:
             skip_parent_only_modules = True
         elif self._children and self.device_type == DeviceType.WallSwitch:
             # _initialize_modules is called on the parent after the children
@@ -209,7 +211,7 @@ class SmartDevice(Device):
             _LOGGER.debug("%s requires %s", mod, mod.REQUIRED_COMPONENT)
 
             if (
-                skip_parent_only_modules and mod in WALL_SWITCH_PARENT_ONLY_MODULES
+                skip_parent_only_modules and mod in NON_HUB_PARENT_ONLY_MODULES
             ) or mod.__name__ in child_modules_to_skip:
                 continue
             if (
diff --git a/kasa/tests/device_fixtures.py b/kasa/tests/device_fixtures.py
index 184eedaa..844314be 100644
--- a/kasa/tests/device_fixtures.py
+++ b/kasa/tests/device_fixtures.py
@@ -152,6 +152,24 @@ def parametrize_combine(parametrized: list[pytest.MarkDecorator]):
     )
 
 
+def parametrize_subtract(params: pytest.MarkDecorator, subtract: pytest.MarkDecorator):
+    """Combine multiple pytest parametrize dev marks into one set of fixtures."""
+    if params.args[0] != "dev" or subtract.args[0] != "dev":
+        raise Exception(
+            f"Supplied mark is not for dev fixture: {params.args[0]} {subtract.args[0]}"
+        )
+    fixtures = []
+    for param in params.args[1]:
+        if param not in subtract.args[1]:
+            fixtures.append(param)
+    return pytest.mark.parametrize(
+        "dev",
+        sorted(fixtures),
+        indirect=True,
+        ids=idgenerator,
+    )
+
+
 def parametrize(
     desc,
     *,
diff --git a/kasa/tests/test_childdevice.py b/kasa/tests/test_childdevice.py
index 9e4b6fdb..26568c24 100644
--- a/kasa/tests/test_childdevice.py
+++ b/kasa/tests/test_childdevice.py
@@ -3,10 +3,20 @@ import sys
 
 import pytest
 
+from kasa.device_type import DeviceType
 from kasa.smart.smartchilddevice import SmartChildDevice
+from kasa.smart.smartdevice import NON_HUB_PARENT_ONLY_MODULES
 from kasa.smartprotocol import _ChildProtocolWrapper
 
-from .conftest import strip_smart
+from .conftest import parametrize, parametrize_subtract, strip_smart
+
+has_children_smart = parametrize(
+    "has children", component_filter="control_child", protocol_filter={"SMART"}
+)
+hub_smart = parametrize(
+    "smart hub", device_type_filter=[DeviceType.Hub], protocol_filter={"SMART"}
+)
+non_hub_parent_smart = parametrize_subtract(has_children_smart, hub_smart)
 
 
 @strip_smart
@@ -82,3 +92,11 @@ async def test_childdevice_properties(dev: SmartChildDevice):
     exceptions = list(_test_property_getters())
     if exceptions:
         raise ExceptionGroup("Accessing child properties caused exceptions", exceptions)
+
+
+@non_hub_parent_smart
+async def test_parent_only_modules(dev, dummy_protocol, mocker):
+    """Test that parent only modules are not available on children."""
+    for child in dev.children:
+        for module in NON_HUB_PARENT_ONLY_MODULES:
+            assert module not in [type(module) for module in child.modules.values()]