Fix zwave_js update entity (#78116)

* Test zwave_js update entity progress

* Block until firmware update is done

* Update homeassistant/components/zwave_js/update.py

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>

* revert params

* unsub finished event listener

* fix tests

* Add test for returned failure

* refactor a little

* rename

* Remove unnecessary controller logic for mocking

* Clear event when resetting

* Comments

* readability

* Fix test

* Fix test

Co-authored-by: Martin Hjelmare <marhje52@gmail.com>
This commit is contained in:
Raman Gupta 2022-09-09 16:10:56 -04:00 committed by GitHub
parent 19cf5dfc6d
commit 8cc0b41daf
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 232 additions and 43 deletions

View File

@ -12,7 +12,12 @@ from zwave_js_server.client import Client as ZwaveClient
from zwave_js_server.const import NodeStatus
from zwave_js_server.exceptions import BaseZwaveJSServerError, FailedZWaveCommand
from zwave_js_server.model.driver import Driver
from zwave_js_server.model.firmware import FirmwareUpdateInfo, FirmwareUpdateProgress
from zwave_js_server.model.firmware import (
FirmwareUpdateFinished,
FirmwareUpdateInfo,
FirmwareUpdateProgress,
FirmwareUpdateStatus,
)
from zwave_js_server.model.node import Node as ZwaveNode
from homeassistant.components.update import UpdateDeviceClass, UpdateEntity
@ -82,7 +87,10 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
self._status_unsub: Callable[[], None] | None = None
self._poll_unsub: Callable[[], None] | None = None
self._progress_unsub: Callable[[], None] | None = None
self._finished_unsub: Callable[[], None] | None = None
self._num_files_installed: int = 0
self._finished_event = asyncio.Event()
self._finished_status: FirmwareUpdateStatus | None = None
# Entity class attributes
self._attr_name = "Firmware"
@ -119,18 +127,38 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
self.async_write_ha_state()
@callback
def _reset_progress(self) -> None:
"""Reset update install progress."""
def _update_finished(self, event: dict[str, Any]) -> None:
"""Update install progress on event."""
finished: FirmwareUpdateFinished = event["firmware_update_finished"]
self._finished_status = finished.status
self._finished_event.set()
@callback
def _unsub_firmware_events_and_reset_progress(
self, write_state: bool = False
) -> None:
"""Unsubscribe from firmware events and reset update install progress."""
if self._progress_unsub:
self._progress_unsub()
self._progress_unsub = None
if self._finished_unsub:
self._finished_unsub()
self._finished_unsub = None
self._finished_status = None
self._finished_event.clear()
self._num_files_installed = 0
self._attr_in_progress = False
self._attr_in_progress = 0
if write_state:
self.async_write_ha_state()
async def _async_update(self, _: HomeAssistant | datetime | None = None) -> None:
"""Update the entity."""
self._poll_unsub = None
# If device is asleep/dead, wait for it to wake up/become alive before
# attempting an update
for status, event_name in (
(NodeStatus.ASLEEP, "wake up"),
(NodeStatus.DEAD, "alive"),
@ -187,19 +215,40 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
"""Install an update."""
firmware = self._latest_version_firmware
assert firmware
self._attr_in_progress = 0
self.async_write_ha_state()
self._unsub_firmware_events_and_reset_progress(True)
self._progress_unsub = self.node.on(
"firmware update progress", self._update_progress
)
self._finished_unsub = self.node.once(
"firmware update finished", self._update_finished
)
for file in firmware.files:
try:
await self.driver.controller.async_begin_ota_firmware_update(
self.node, file
)
except BaseZwaveJSServerError as err:
self._reset_progress()
self._unsub_firmware_events_and_reset_progress()
raise HomeAssistantError(err) from err
# We need to block until we receive the `firmware update finished` event
await self._finished_event.wait()
assert self._finished_status is not None
# If status is not OK, we should throw an error to let the user know
if self._finished_status not in (
FirmwareUpdateStatus.OK_NO_RESTART,
FirmwareUpdateStatus.OK_RESTART_PENDING,
FirmwareUpdateStatus.OK_WAITING_FOR_ACTIVATION,
):
status = self._finished_status
self._unsub_firmware_events_and_reset_progress()
raise HomeAssistantError(status.name.replace("_", " ").title())
# If we get here, the firmware installation was successful and we need to
# update progress accordingly
self._num_files_installed += 1
self._attr_in_progress = floor(
100 * self._num_files_installed / len(firmware.files)
@ -208,7 +257,7 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
self._attr_installed_version = self._attr_latest_version = firmware.version
self._latest_version_firmware = None
self._reset_progress()
self._unsub_firmware_events_and_reset_progress()
async def async_poll_value(self, _: bool) -> None:
"""Poll a value."""
@ -255,6 +304,4 @@ class ZWaveNodeFirmwareUpdate(UpdateEntity):
self._poll_unsub()
self._poll_unsub = None
if self._progress_unsub:
self._progress_unsub()
self._progress_unsub = None
self._unsub_firmware_events_and_reset_progress()

View File

@ -1,9 +1,11 @@
"""Test the Z-Wave JS update entities."""
import asyncio
from datetime import timedelta
import pytest
from zwave_js_server.event import Event
from zwave_js_server.exceptions import FailedZWaveCommand
from zwave_js_server.model.firmware import FirmwareUpdateStatus
from homeassistant.components.update.const import (
ATTR_AUTO_UPDATE,
@ -51,7 +53,7 @@ FIRMWARE_UPDATES = {
}
async def test_update_entity_success(
async def test_update_entity_states(
hass,
client,
climate_radio_thermostat_ct100_plus_different_endpoints,
@ -60,7 +62,7 @@ async def test_update_entity_success(
caplog,
hass_ws_client,
):
"""Test update entity."""
"""Test update entity states."""
ws_client = await hass_ws_client(hass)
await hass.async_block_till_done()
@ -137,39 +139,14 @@ async def test_update_entity_success(
client.async_send_command.reset_mock()
# Test successful install call without a version
await hass.services.async_call(
UPDATE_DOMAIN,
SERVICE_INSTALL,
{
ATTR_ENTITY_ID: UPDATE_ENTITY,
},
blocking=True,
)
args = client.async_send_command.call_args_list[0][0][0]
assert args["command"] == "controller.begin_ota_firmware_update"
assert (
args["nodeId"]
== climate_radio_thermostat_ct100_plus_different_endpoints.node_id
)
assert args["update"] == {
"target": 0,
"url": "https://example2.com",
"integrity": "sha2",
}
client.async_send_command.reset_mock()
async def test_update_entity_install_failure(
async def test_update_entity_install_raises(
hass,
client,
climate_radio_thermostat_ct100_plus_different_endpoints,
controller_node,
integration,
):
"""Test update entity failed install."""
"""Test update entity install raises exception."""
client.async_send_command.return_value = FIRMWARE_UPDATES
async_fire_time_changed(hass, dt_util.utcnow() + timedelta(days=1))
@ -287,11 +264,10 @@ async def test_update_entity_ha_not_running(
assert args["nodeId"] == zen_31.node_id
async def test_update_entity_failure(
async def test_update_entity_update_failure(
hass,
client,
climate_radio_thermostat_ct100_plus_different_endpoints,
controller_node,
integration,
):
"""Test update entity update failed."""
@ -311,3 +287,169 @@ async def test_update_entity_failure(
args["nodeId"]
== climate_radio_thermostat_ct100_plus_different_endpoints.node_id
)
async def test_update_entity_progress(
hass,
client,
climate_radio_thermostat_ct100_plus_different_endpoints,
integration,
):
"""Test update entity progress."""
node = climate_radio_thermostat_ct100_plus_different_endpoints
client.async_send_command.return_value = FIRMWARE_UPDATES
async_fire_time_changed(hass, dt_util.utcnow() + timedelta(days=1))
await hass.async_block_till_done()
state = hass.states.get(UPDATE_ENTITY)
assert state
assert state.state == STATE_ON
attrs = state.attributes
assert attrs[ATTR_INSTALLED_VERSION] == "10.7"
assert attrs[ATTR_LATEST_VERSION] == "11.2.4"
client.async_send_command.reset_mock()
client.async_send_command.return_value = None
# Test successful install call without a version
install_task = hass.async_create_task(
hass.services.async_call(
UPDATE_DOMAIN,
SERVICE_INSTALL,
{
ATTR_ENTITY_ID: UPDATE_ENTITY,
},
blocking=True,
)
)
# Sleep so that task starts
await asyncio.sleep(0.1)
event = Event(
type="firmware update progress",
data={
"source": "node",
"event": "firmware update progress",
"nodeId": node.node_id,
"sentFragments": 1,
"totalFragments": 20,
},
)
node.receive_event(event)
# Validate that the progress is updated
state = hass.states.get(UPDATE_ENTITY)
assert state
attrs = state.attributes
assert attrs[ATTR_IN_PROGRESS] == 5
event = Event(
type="firmware update finished",
data={
"source": "node",
"event": "firmware update finished",
"nodeId": node.node_id,
"status": FirmwareUpdateStatus.OK_NO_RESTART,
},
)
node.receive_event(event)
await hass.async_block_till_done()
# Validate that progress is reset and entity reflects new version
state = hass.states.get(UPDATE_ENTITY)
assert state
attrs = state.attributes
assert attrs[ATTR_IN_PROGRESS] is False
assert attrs[ATTR_INSTALLED_VERSION] == "11.2.4"
assert attrs[ATTR_LATEST_VERSION] == "11.2.4"
assert state.state == STATE_OFF
await install_task
async def test_update_entity_install_failed(
hass,
client,
climate_radio_thermostat_ct100_plus_different_endpoints,
integration,
caplog,
):
"""Test update entity install returns error status."""
node = climate_radio_thermostat_ct100_plus_different_endpoints
client.async_send_command.return_value = FIRMWARE_UPDATES
async_fire_time_changed(hass, dt_util.utcnow() + timedelta(days=1))
await hass.async_block_till_done()
state = hass.states.get(UPDATE_ENTITY)
assert state
assert state.state == STATE_ON
attrs = state.attributes
assert attrs[ATTR_INSTALLED_VERSION] == "10.7"
assert attrs[ATTR_LATEST_VERSION] == "11.2.4"
client.async_send_command.reset_mock()
client.async_send_command.return_value = None
async def call_install():
await hass.services.async_call(
UPDATE_DOMAIN,
SERVICE_INSTALL,
{
ATTR_ENTITY_ID: UPDATE_ENTITY,
},
blocking=True,
)
# Test install call - we expect it to raise
install_task = hass.async_create_task(call_install())
# Sleep so that task starts
await asyncio.sleep(0.1)
event = Event(
type="firmware update progress",
data={
"source": "node",
"event": "firmware update progress",
"nodeId": node.node_id,
"sentFragments": 1,
"totalFragments": 20,
},
)
node.receive_event(event)
# Validate that the progress is updated
state = hass.states.get(UPDATE_ENTITY)
assert state
attrs = state.attributes
assert attrs[ATTR_IN_PROGRESS] == 5
event = Event(
type="firmware update finished",
data={
"source": "node",
"event": "firmware update finished",
"nodeId": node.node_id,
"status": FirmwareUpdateStatus.ERROR_TIMEOUT,
},
)
node.receive_event(event)
await hass.async_block_till_done()
# Validate that progress is reset and entity reflects old version
state = hass.states.get(UPDATE_ENTITY)
assert state
attrs = state.attributes
assert attrs[ATTR_IN_PROGRESS] is False
assert attrs[ATTR_INSTALLED_VERSION] == "10.7"
assert attrs[ATTR_LATEST_VERSION] == "11.2.4"
assert state.state == STATE_ON
# validate that the install task failed
with pytest.raises(HomeAssistantError):
await install_task