Migrate RuleModule to mashumaro (#1283)

Also fixes a bug whereby multiple queries for the same module would overwrite each other.

kasa/iot/iotdevice.py

@@ -412,7 +412,15 @@ class IotDevice(Device):
         # every other update will query for them
         update: dict = self._last_update.copy() if self._last_update else {}
         for response in responses:
-            update = {**update, **response}
+            for k, v in response.items():
+                # The same module could have results in different responses
+                # i.e. smartlife.iot.common.schedule for Usage and
+                # Schedule, so need to call update(**v) here. If a module is
+                # not supported the response
+                # {'err_code': -1, 'err_msg': 'module not support'}
+                # become top level key/values of the response so check for dict
+                if isinstance(v, dict):
+                    update.setdefault(k, {}).update(**v)
         self._last_update = update
 
         # IOT modules are added as default but could be unsupported post first update

kasa/iot/modules/rulemodule.py

@@ -3,9 +3,10 @@
 from __future__ import annotations
 
 import logging
+from dataclasses import dataclass
 from enum import Enum
 
-from pydantic.v1 import BaseModel
+from mashumaro import DataClassDictMixin
 
 from ..iotmodule import IotModule, merge
 
@@ -28,26 +29,27 @@ class TimeOption(Enum):
     AtSunset = 2
 
 
-class Rule(BaseModel):
+@dataclass
+class Rule(DataClassDictMixin):
     """Representation of a rule."""
 
     id: str
     name: str
-    enable: bool
+    enable: int
     wday: list[int]
-    repeat: bool
+    repeat: int
 
     # start action
-    sact: Action | None
+    sact: Action | None = None
-    stime_opt: TimeOption
+    stime_opt: TimeOption | None = None
-    smin: int
+    smin: int | None = None
 
-    eact: Action | None
+    eact: Action | None = None
-    etime_opt: TimeOption
+    etime_opt: TimeOption | None = None
-    emin: int
+    emin: int | None = None
 
     # Only on bulbs
-    s_light: dict | None
+    s_light: dict | None = None
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -66,7 +68,7 @@ class RuleModule(IotModule):
         """Return the list of rules for the service."""
         try:
             return [
-                Rule.parse_obj(rule) for rule in self.data["get_rules"]["rule_list"]
+                Rule.from_dict(rule) for rule in self.data["get_rules"]["rule_list"]
             ]
         except Exception as ex:
             _LOGGER.error("Unable to read rule list: %s (data: %s)", ex, self.data)

tests/fakeprotocol_iot.py

@@ -136,6 +136,34 @@ CLOUD_MODULE = {
     }
 }
 
+SCHEDULE_MODULE = {
+    "get_next_action": {
+        "action": 1,
+        "err_code": 0,
+        "id": "0794F4729DB271627D1CF35A9A854030",
+        "schd_time": 68927,
+        "type": 2,
+    },
+    "get_rules": {
+        "enable": 1,
+        "err_code": 0,
+        "rule_list": [
+            {
+                "eact": -1,
+                "enable": 1,
+                "id": "8AA75A50A8440B17941D192BD9E01FFA",
+                "name": "name",
+                "repeat": 1,
+                "sact": 1,
+                "smin": 1027,
+                "soffset": 0,
+                "stime_opt": 2,
+                "wday": [1, 1, 1, 1, 1, 1, 1],
+            },
+        ],
+        "version": 2,
+    },
+}
 
 AMBIENT_MODULE = {
     "get_current_brt": {"value": 26, "err_code": 0},
@@ -450,6 +478,8 @@ class FakeIotTransport(BaseTransport):
         "smartlife.iot.PIR": MOTION_MODULE,
         "cnCloud": CLOUD_MODULE,
         "smartlife.iot.common.cloud": CLOUD_MODULE,
+        "schedule": SCHEDULE_MODULE,
+        "smartlife.iot.common.schedule": SCHEDULE_MODULE,
     }
 
     async def send(self, request, port=9999):

tests/iot/modules/test_schedule.py

@@ -0,0 +1,17 @@
+import pytest
+
+from kasa import Device, Module
+from kasa.iot.modules.rulemodule import Action, TimeOption
+
+from ...device_fixtures import device_iot
+
+
+@device_iot
+def test_schedule(dev: Device, caplog: pytest.LogCaptureFixture):
+    schedule = dev.modules.get(Module.IotSchedule)
+    assert schedule
+    if rules := schedule.rules:
+        first = rules[0]
+        assert isinstance(first.sact, Action)
+        assert isinstance(first.stime_opt, TimeOption)
+    assert "Unable to read rule list" not in caplog.text