8000 Fix update logic & add tests · python-kasa/python-kasa@da83cd5 · GitHub
[go: up one dir, main page]

Skip to content

Commit da83cd5

Browse files
committed
Fix update logic & add tests
1 parent 39e6aac commit da83cd5

File tree

3 files changed

+156
-14
lines changed

3 files changed

+156
-14
lines changed

kasa/smart/modules/firmware.py

Lines changed: 53 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,19 @@
2323
_LOGGER = logging.getLogger(__name__)
2424

2525

26+
class DownloadState(BaseModel):
27+
"""Download state."""
28+
29+
# Example:
30+
# {'status': 0, 'download_progress': 0, 'reboot_time': 5,
31+
# 'upgrade_time': 5, 'auto_upgrade': False}
32+
status: int
33+
progress: int = Field(alias="download_progress")
34+
reboot_time: int
35+
upgrade_time: int
36+
auto_upgrade: bool
37+
38+
2639
class UpdateInfo(BaseModel):
2740
"""Update info status object."""
2841

@@ -86,7 +99,7 @@ def __init__(self, device: SmartDevice, module: str):
8699
name="Current firmware version",
87100
container=self,
88101
attribute_getter="current_firmware",
89-
category=Feature.Category.Info,
102+
category=Feature.Category.Debug,
90103
)
91104
)
92105
self._add_feature(
@@ -96,7 +109,7 @@ def __init__(self, device: SmartDevice, module: str):
96109
name="Available firmware version",
97110
container=self,
98111
attribute_getter="latest_firmware",
99-
category=Feature.Category.Info,
112+
category=Feature.Category.Debug,
100113
)
101114
)
102115

@@ -134,31 +147,58 @@ def update_available(self) -> bool | None:
134147
return None
135148
return self.firmware_update_info.update_available
136149

137-
async def get_update_state(self):
150+
async def get_update_state(self) -> DownloadState:
138151
"""Return update state."""
139-
return await self.call("get_fw_download_state")
152+
resp = await self.call("get_fw_download_state")
153+
state = resp["get_fw_download_state"]
154+
return DownloadState(**state)
140155

141-
async def update(self):
156+
async def update(self, progress_cb=None):
142157
"""Update the device firmware."""
143158
current_fw = self.current_firmware
144-
_LOGGER.debug(
159+
_LOGGER.info(
145160
"Going to upgrade from %s to %s",
146161
current_fw,
147162
self.firmware_update_info.version,
148163
)
149-
resp = await self.call("fw_download")
150-
_LOGGER.debug("Update request response: %s", resp)
164+
await self.call("fw_download")
165+
151166
# TODO: read timeout from get_auto_update_info or from get_fw_download_state?
152167
async with asyncio_timeout(60 * 5):
153168
while True:
154169
await asyncio.sleep(0.5)
155-
state = await self.get_update_state()
156-
_LOGGER.debug("Update state: %s" % state)
157-
# TODO: this could await a given callable for progress
170+
try:
171+
state = await self.get_update_state()
172+
except Exception as ex:
173+
_LOGGER.warning(
174+
"Got exception, maybe the device is rebooting? %s", ex
175+
)
176+
continue
158177

159-
if self.firmware_update_info.version != current_fw:
160-
_LOGGER.info("Updated to %s", self.firmware_update_info.version)
178+
_LOGGER.debug("Update state: %s" % state)
179+
if progress_cb is not None:
180+
asyncio.ensure_future(progress_cb(state))
181+
182+
if state.status == 0:
183+
_LOGGER.info(
184+
"Update idle, hopefully updated to %s",
185+
self.firmware_update_info.version,
186+
)
187+
break
188+
elif state.status == 2:
189+
_LOGGER.info("Downloading firmware, progress: %s", state.progress)
190+
elif state.status == 3:
191+
upgrade_sleep = state.upgrade_time
192+
_LOGGER.info(
193+
"Flashing firmware, sleeping for %s before checking status",
194+
upgrade_sleep,
195+
)
196+
await asyncio.sleep(upgrade_sleep)
197+
elif state.status < 0:
198+
_LOGGER.error("Got error: %s", state.status)
161199
break
200+
else:
201+
_LOGGER.warning("Unhandled state code: %s", state)
162202

163203
@property
164204
def auto_update_enabled(self):

kasa/tests/fakeprotocol_smart.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -234,7 +234,7 @@ def _send_request(self, request_dict: dict):
234234
pytest.fixtures_missing_methods[self.fixture_name] = set()
235235
pytest.fixtures_missing_methods[self.fixture_name].add(method)
236236
return retval
237-
elif method == "set_qs_info":
237+
elif method in ["set_qs_info", "fw_download"]:
238238
return {"error_code": 0}
239239
elif method == "set_dynamic_light_effect_rule_enable":
240240
self._set_light_effect(info, params)
Lines changed: 102 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,102 @@
1+
from __future__ import annotations
2+
3+
import logging
4+
5+
import pytest
6+
from pytest_mock import MockerFixture
7+
8+
from kasa.smart import SmartDevice
9+
from kasa.smart.modules import Firmware
10+
from kasa.smart.modules.firmware import DownloadState
11+
from kasa.tests.device_fixtures import parametrize
12+
13+
firmware = parametrize(
14+
"has firmware", component_filter="firmware", protocol_filter={"SMART"}
15+
)
16+
17+
18+
@firmware
19+
@pytest.mark.parametrize(
20+
"feature, prop_name, type, required_version",
21+
[
22+
("auto_update_enabled", "auto_update_enabled", bool, 2),
23+
("update_available", "update_available", bool, 1),
24+
("update_available", "update_available", bool, 1),
25+
("current_firmware_version", "current_firmware", str, 1),
26+
("available_firmware_version", "latest_firmware", str, 1),
27+
],
28+
)
29+
async def test_firmware_features(
30+
dev: SmartDevice, feature, prop_name, type, required_version, mocker: MockerFixture
31+
):
32+
"""Test light effect."""
33+
fw = dev.get_module(Firmware)
34+
assert fw
35+
36+
if not dev.is_cloud_connected:
37+
pytest.skip("Device is not cloud connected, skipping test")
38+
39+
if fw.supported_version < required_version:
40+
pytest.skip("Feature %s requires newer version" % feature)
41+
42+
prop = getattr(fw, prop_name)
43+
assert isinstance(prop, type)
44+
45+
feat = fw._module_features[feature]
46+
assert feat.value == prop
47+
assert isinstance(feat.value, type)
48+
49+
50+
@firmware
51+
async def test_update_available_without_cloud(dev: SmartDevice):
52+
"""Test that update_available returns None when disconnected."""
53+
fw = dev.get_module(Firmware)
54+
assert fw
55+
56+
if dev.is_cloud_connected:
57+
assert isinstance(fw.update_available, bool)
58+
else:
59+
assert fw.update_available is None
60+
61+
62+
@firmware
63+
async def test_update(
64+
dev: SmartDevice, mocker: MockerFixture, caplog: pytest.LogCaptureFixture
65+
):
66+
"""Test updating firmware."""
67+
caplog.set_level(logging.INFO)
68+
69+
fw = dev.get_module(Firmware)
70+
assert fw
71+
72+
upgrade_time = 5
73+
extras = {"reboot_time": 5, "upgrade_time": upgrade_time, "auto_upgrade": False}
74+
update_states = [
75+
# Unknown 1
76+
DownloadState(status=1, download_progress=0, **extras),
77+
# Downloading
78+
DownloadState(status=2, download_progress=10, **extras),
79+
DownloadState(status=2, download_progress=100, **extras),
80+
# Flashing
81+
DownloadState(status=3, download_progress=100, **extras),
82+
DownloadState(status=3, download_progress=100, **extras),
83+
# Done
84+
DownloadState(status=0, download_progress=100, **extras),
85+
]
86+
87+
sleep = mocker.patch("asyncio.sleep")
88+
mocker.patch.object(fw, "get_update_state", side_effect=update_states)
89+
90+
cb_mock = mocker.AsyncMock()
91+
92+
await fw.update(progress_cb=cb_mock)
93+
94+
assert "Unhandled state code" in caplog.text
95+
assert "Downloading firmware, progress: 10" in caplog.text
96+
assert "Flashing firmware, sleeping" in caplog.text
97+
assert "Update idle" in caplog.text
98+
99+
cb_mock.assert_called()
100+
101+
# sleep based on the upgrade_time
102+
sleep.assert_any_call(upgrade_time)

0 commit comments

Comments
 (0)
0