8000 Migrate smart firmware module to mashumaro (#1276) · ryenitcher/python-kasa@999e84d · GitHub
[go: up one dir, main page]
Skip to content

Commit 999e84d

Browse files
sdb9696sdb9696
authored
Migrate smart firmware module to mashumaro (python-kasa#1276)
1 parent 03c073c commit 999e84d

File tree

2 files changed

+43
-26
lines changed

2 files changed

+43
-26
lines changed

‎kasa/smart/modules/firmware.py

Lines changed: 22 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -6,10 +6,12 @@
66
import logging
77
from asyncio import timeout as asyncio_timeout
88
from collections.abc import Callable, Coroutine
9+
from dataclasses import dataclass, field
910
from datetime import date
10-
from typing import TYPE_CHECKING
11+
from typing import TYPE_CHECKING, Annotated
1112

12-
from pydantic.v1 import BaseModel, Field, validator
13+
from mashumaro import DataClassDictMixin, field_options
14+
from mashumaro.types import Alias
1315

1416
from ...exceptions import KasaException
1517
from ...feature import Feature
@@ -22,36 +24,36 @@
2224
_LOGGER = logging.getLogger(__name__)
2325

2426

25-
class DownloadState(BaseModel):
27+
@dataclass
28+
class DownloadState(DataClassDictMixin):
2629
"""Download state."""
2730

2831
# Example:
2932
# {'status': 0, 'download_progress': 0, 'reboot_time': 5,
3033
# 'upgrade_time': 5, 'auto_upgrade': False}
3134
status: int
32-
progress: int = Field(alias="download_progress")
35+
progress: Annotated[int, Alias("download_progress")]
3336
reboot_time: int
3437
upgrade_time: int
3538
auto_upgrade: bool
3639

3740

38-
class UpdateInfo(BaseModel):
41+
@dataclass
42+
class UpdateInfo(DataClassDictMixin):
3943
"""Update info status object."""
4044

41-
status: int = Field(alias="type")
42-
version: str | None = Field(alias="fw_ver", default=None)
43-
release_date: date | None = None
44-
release_notes: str | None = Field(alias="release_note", default=None)
45+
status: Annotated[int, Alias("type")]
46+
needs_upgrade: Annotated[bool, Alias("need_to_upgrade")]
47+
version: Annotated[str | None, Alias("fw_ver")] = None
48+
release_date: date | None = field(
49+
default=None,
50+
metadata=field_options(
51+
deserialize=lambda x: date.fromisoformat(x) if x else None
52+
),
53+
)
54+
release_notes: Annotated[str | None, Alias("release_note")] = None
4555
fw_size: int | None = None
4656
oem_id: str | None = None
47-
needs_upgrade: bool = Field(alias="need_to_upgrade")
48-
49-
@validator("release_date", pre=True)
50-
def _release_date_optional(cls, v: str) -> str | None:
51-
if not v:
52-
return None
53-
54-
return v
5557

5658
@property
5759
def update_available(self) -> bool:
@@ -139,7 +141,7 @@ async def check_latest_firmware(self) -> UpdateInfo | None:
139141
"""Check for the latest firmware for the device."""
140142
try:
141143
fw = await self.call("get_latest_fw")
142-
self._firmware_update_info = UpdateInfo.parse_obj(fw["get_latest_fw"])
144+
self._firmware_update_info = UpdateInfo.from_dict(fw["get_latest_fw"])
143145
return self._firmware_update_info
144146
except Exception:
145147
_LOGGER.exception("Error getting latest firmware for %s:", self._device)
@@ -174,7 +176,7 @@ async def get_update_state(self) -> DownloadState:
174176
"""Return update state."""
175177
resp = await self.call("get_fw_download_state")
176178
state = resp["get_fw_download_state"]
177-
return DownloadState(**state)
179+
return DownloadState.from_dict(state)
178180

179181
@allow_update_after
180182
async def update(
@@ -232,7 +234,7 @@ async def update(
232234
else:
233235
_LOGGER.warning("Unhandled state code: %s", state)
234236

235-
return state.dict()
237+
return state.to_dict()
236238

237239
@property
238240
def auto_update_enabled(self) -> bool:

‎tests/smart/modules/test_firmware.py

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
import asyncio
44
import logging
55
from contextlib import nullcontext
6+
from datetime import date
67
from typing import TypedDict
78

89
import pytest
@@ -52,6 +53,20 @@ async def test_firmware_features(
5253
assert isinstance(feat.value, type)
5354

5455

56+
@firmware
57+
async def test_firmware_update_info(dev: SmartDevice):
58+
"""Test that the firmware UpdateInfo object deserializes correctly."""
59+
fw = dev.modules.get(Module.Firmware)
60+
assert fw
61+
62+
if not dev.is_cloud_connected:
63+
pytest.skip("Device is not cloud connected, skipping test")
64+
assert fw.firmware_update_info is None
65+
await fw.check_latest_firmware()
66+
assert fw.firmware_update_info is not None
67+
assert isinstance(fw.firmware_update_info.release_date, date | None)
68+
69+
5570
@firmware
5671
async def test_update_available_without_cloud(dev: SmartDevice):
5772
"""Test that update_available returns None when disconnected."""
@@ -105,15 +120,15 @@ class Extras(TypedDict):
105120
}
106121
update_states = [
107122
# Unknown 1
108-
DownloadState(status=1, download_progress=0, **extras),
123+
DownloadState(status=1, progress=0, **extras),
109124
# Downloading
110-
DownloadState(status=2, download_progress=10, **extras),
111-
DownloadState(status=2, download_progress=100, **extras),
125+
DownloadState(status=2, progress=10, **extras),
126+
DownloadState(status=2, progress=100, **extras),
112127
# Flashing
113-
DownloadState(status=3, download_progress=100, **extras),
114-
DownloadState(status=3, download_progress=100, **extras),
128+
DownloadState(status=3, progress=100, **extras),
129+
DownloadState(status=3, progress=100, **extras),
115130
# Done
116-
DownloadState(status=0, download_progress=100, **extras),
131+
DownloadState(status=0, progress=100, **extras),
117132
]
118133

119134
asyncio_sleep = asyncio.sleep

0 commit comments

Comments
 (0)
0