From 058c0c7db1aa08bd1d9a5474a14e90c4279ddab3 Mon Sep 17 00:00:00 2001
From: Steven B <51370195+sdb9696@users.noreply.github.com>
Date: Wed, 20 Nov 2024 13:17:26 +0000
Subject: [PATCH] Migrate TurnOnBehaviours to mashumaro

---
 kasa/iot/iotbulb.py       | 52 ++++++++++++++++++++-------------------
 tests/fakeprotocol_iot.py | 21 ++++++++++++++++
 tests/test_bulb.py        |  6 +++++
 3 files changed, 54 insertions(+), 25 deletions(-)

diff --git a/kasa/iot/iotbulb.py b/kasa/iot/iotbulb.py
index 14c711031..cb2e858cd 100644
--- a/kasa/iot/iotbulb.py
+++ b/kasa/iot/iotbulb.py
@@ -4,10 +4,13 @@
 
 import logging
 import re
+from dataclasses import dataclass
 from enum import Enum
-from typing import cast
+from typing import Annotated, cast
 
-from pydantic.v1 import BaseModel, Field, root_validator
+from mashumaro import DataClassDictMixin
+from mashumaro.config import BaseConfig
+from mashumaro.types import Alias
 
 from ..device_type import DeviceType
 from ..deviceconfig import DeviceConfig
@@ -35,9 +38,12 @@ class BehaviorMode(str, Enum):
     Last = "last_status"
     #: Use chosen preset.
     Preset = "customize_preset"
+    #: Circadian
+    Circadian = "circadian"
 
 
-class TurnOnBehavior(BaseModel):
+@dataclass
+class TurnOnBehavior(DataClassDictMixin):
     """Model to present a single turn on behavior.
 
     :param int preset: the index number of wanted preset.
@@ -48,34 +54,30 @@ class TurnOnBehavior(BaseModel):
     to contain either the preset index, or ``None`` for the last known state.
     """
 
-    #: Index of preset to use, or ``None`` for the last known state.
-    preset: int | None = Field(alias="index", default=None)
-    #: Wanted behavior
-    mode: BehaviorMode
-
-    @root_validator
-    def _mode_based_on_preset(cls, values: dict) -> dict:
-        """Set the mode based on the preset value."""
-        if values["preset"] is not None:
-            values["mode"] = BehaviorMode.Preset
-        else:
-            values["mode"] = BehaviorMode.Last
+    class Config(BaseConfig):
+        """Serialization config."""
 
-        return values
+        omit_none = True
+        serialize_by_alias = True
 
-    class Config:
-        """Configuration to make the validator run when changing the values."""
-
-        validate_assignment = True
+    #: Wanted behavior
+    mode: BehaviorMode
+    #: Index of preset to use, or ``None`` for the last known state.
+    preset: Annotated[int | None, Alias("index")] = None
+    brightness: int | None = None
+    color_temp: int | None = None
+    hue: int | None = None
+    saturation: int | None = None
 
 
-class TurnOnBehaviors(BaseModel):
+@dataclass
+class TurnOnBehaviors(DataClassDictMixin):
     """Model to contain turn on behaviors."""
 
     #: The behavior when the bulb is turned on programmatically.
-    soft: TurnOnBehavior = Field(alias="soft_on")
+    soft: Annotated[TurnOnBehavior, Alias("soft_on")]
     #: The behavior when the bulb has been off from mains power.
-    hard: TurnOnBehavior = Field(alias="hard_on")
+    hard: Annotated[TurnOnBehavior, Alias("hard_on")]
 
 
 TPLINK_KELVIN = {
@@ -303,7 +305,7 @@ async def get_light_details(self) -> dict[str, int]:
 
     async def get_turn_on_behavior(self) -> TurnOnBehaviors:
         """Return the behavior for turning the bulb on."""
-        return TurnOnBehaviors.parse_obj(
+        return TurnOnBehaviors.from_dict(
             await self._query_helper(self.LIGHT_SERVICE, "get_default_behavior")
         )
 
@@ -314,7 +316,7 @@ async def set_turn_on_behavior(self, behavior: TurnOnBehaviors) -> dict:
         you should use :func:`get_turn_on_behavior` to get the current settings.
         """
         return await self._query_helper(
-            self.LIGHT_SERVICE, "set_default_behavior", behavior.dict(by_alias=True)
+            self.LIGHT_SERVICE, "set_default_behavior", behavior.to_dict()
         )
 
     async def get_light_state(self) -> dict[str, dict]:
diff --git a/tests/fakeprotocol_iot.py b/tests/fakeprotocol_iot.py
index 8c4e4057c..9a236813d 100644
--- a/tests/fakeprotocol_iot.py
+++ b/tests/fakeprotocol_iot.py
@@ -175,6 +175,23 @@ def success(res):
     }
 }
 
+LIGHT_DETAILS = {
+    "color_rendering_index": 80,
+    "err_code": 0,
+    "incandescent_equivalent": 60,
+    "lamp_beam_angle": 150,
+    "max_lumens": 800,
+    "max_voltage": 120,
+    "min_voltage": 110,
+    "wattage": 10,
+}
+
+DEFAULT_BEHAVIOR = {
+    "err_code": 0,
+    "hard_on": {"mode": "circadian"},
+    "soft_on": {"mode": "last_status"},
+}
+
 
 class FakeIotProtocol(IotProtocol):
     def __init__(self, info, fixture_name=None, *, verbatim=False):
@@ -398,6 +415,8 @@ def set_time(self, new_state: dict, *args):
         },
         "smartlife.iot.smartbulb.lightingservice": {
             "get_light_state": light_state,
+            "get_light_details": LIGHT_DETAILS,
+            "get_default_behavior": DEFAULT_BEHAVIOR,
             "transition_light_state": transition_light_state,
             "set_preferred_state": set_preferred_state,
         },
@@ -408,6 +427,8 @@ def set_time(self, new_state: dict, *args):
         "smartlife.iot.lightStrip": {
             "set_light_state": transition_light_state,
             "get_light_state": light_state,
+            "get_light_details": LIGHT_DETAILS,
+            "get_default_behavior": DEFAULT_BEHAVIOR,
             "set_preferred_state": set_preferred_state,
         },
         "smartlife.iot.common.system": {
diff --git a/tests/test_bulb.py b/tests/test_bulb.py
index ac4400731..f74ed8ec0 100644
--- a/tests/test_bulb.py
+++ b/tests/test_bulb.py
@@ -494,3 +494,9 @@ async def test_modify_preset_payloads(dev: IotBulb, preset, payload, mocker):
 @bulb
 def test_device_type_bulb(dev: Device):
     assert dev.device_type in {DeviceType.Bulb, DeviceType.LightStrip}
+
+
+@bulb_iot
+async def test_turn_on_behaviours(dev: IotBulb):
+    behavior = await dev.get_turn_on_behavior()
+    assert behavior