From 5f4103a4a7d7343751417c5428f7e817ad20c831 Mon Sep 17 00:00:00 2001
From: TimL <tl@smlight.tech>
Date: Wed, 30 Oct 2024 18:02:30 +1100
Subject: [PATCH] Allow smlight device to reboot before updating firmware data
 coordinator (#127442)

* Add delay before updating firmware coordinator

* fix update tests

* change sleep to 1s

* Timeout incase reboot fails

* update test

* test reboot timeout

* log hostname in warning
---
 homeassistant/components/smlight/update.py | 21 +++++++--
 tests/components/smlight/test_update.py    | 55 +++++++++++++++++++++-
 2 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/smlight/update.py b/homeassistant/components/smlight/update.py
index c1149fe3315..147b1d766ef 100644
--- a/homeassistant/components/smlight/update.py
+++ b/homeassistant/components/smlight/update.py
@@ -23,6 +23,7 @@ from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from . import SmConfigEntry
+from .const import LOGGER
 from .coordinator import SmFirmwareUpdateCoordinator, SmFwData
 from .entity import SmEntity
 
@@ -159,7 +160,6 @@ class SmUpdateEntity(SmEntity, UpdateEntity):
     def _update_done(self) -> None:
         """Handle cleanup for update done."""
         self._finished_event.set()
-        self.coordinator.in_progress = False
 
         for remove_cb in self._unload:
             remove_cb()
@@ -178,7 +178,7 @@ class SmUpdateEntity(SmEntity, UpdateEntity):
     @callback
     def _update_failed(self, event: MessageEvent) -> None:
         self._update_done()
-
+        self.coordinator.in_progress = False
         raise HomeAssistantError(f"Update failed for {self.name}")
 
     async def async_install(
@@ -197,5 +197,20 @@ class SmUpdateEntity(SmEntity, UpdateEntity):
             # block until update finished event received
             await self._finished_event.wait()
 
-            await self.coordinator.async_refresh()
+            # allow time for SLZB-06 to reboot before updating coordinator data
+            try:
+                async with asyncio.timeout(180):
+                    while (
+                        self.coordinator.in_progress
+                        and self.installed_version != self._firmware.ver
+                    ):
+                        await self.coordinator.async_refresh()
+                        await asyncio.sleep(1)
+            except TimeoutError:
+                LOGGER.warning(
+                    "Timeout waiting for %s to reboot after update",
+                    self.coordinator.data.info.hostname,
+                )
+
+            self.coordinator.in_progress = False
             self._finished_event.clear()
diff --git a/tests/components/smlight/test_update.py b/tests/components/smlight/test_update.py
index 714caefd91c..0bb2e34d7ca 100644
--- a/tests/components/smlight/test_update.py
+++ b/tests/components/smlight/test_update.py
@@ -1,6 +1,7 @@
 """Tests for the SMLIGHT update platform."""
 
-from unittest.mock import MagicMock
+from datetime import timedelta
+from unittest.mock import MagicMock, patch
 
 from freezegun.api import FrozenDateTimeFactory
 from pysmlight import Firmware, Info
@@ -88,7 +89,9 @@ async def test_update_setup(
     await hass.config_entries.async_unload(entry.entry_id)
 
 
+@patch("homeassistant.components.smlight.update.asyncio.sleep", return_value=None)
 async def test_update_firmware(
+    mock_sleep: MagicMock,
     hass: HomeAssistant,
     freezer: FrozenDateTimeFactory,
     mock_config_entry: MockConfigEntry,
@@ -126,7 +129,7 @@ async def test_update_firmware(
         sw_version="v2.5.2",
     )
 
-    freezer.tick(SCAN_FIRMWARE_INTERVAL)
+    freezer.tick(timedelta(seconds=5))
     async_fire_time_changed(hass)
     await hass.async_block_till_done()
 
@@ -216,6 +219,54 @@ async def test_update_firmware_failed(
     assert state.attributes[ATTR_UPDATE_PERCENTAGE] is None
 
 
+@patch("homeassistant.components.smlight.const.LOGGER.warning")
+async def test_update_reboot_timeout(
+    mock_warning: MagicMock,
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    mock_config_entry: MockConfigEntry,
+    mock_smlight_client: MagicMock,
+) -> None:
+    """Test firmware updates."""
+    await setup_integration(hass, mock_config_entry)
+    entity_id = "update.mock_title_core_firmware"
+    state = hass.states.get(entity_id)
+    assert state.state == STATE_ON
+    assert state.attributes[ATTR_INSTALLED_VERSION] == "v2.3.6"
+    assert state.attributes[ATTR_LATEST_VERSION] == "v2.5.2"
+
+    with (
+        patch(
+            "homeassistant.components.smlight.update.asyncio.timeout",
+            side_effect=TimeoutError,
+        ),
+        patch(
+            "homeassistant.components.smlight.update.asyncio.sleep",
+            return_value=None,
+        ),
+    ):
+        await hass.services.async_call(
+            PLATFORM,
+            SERVICE_INSTALL,
+            {ATTR_ENTITY_ID: entity_id},
+            blocking=False,
+        )
+
+        assert len(mock_smlight_client.fw_update.mock_calls) == 1
+
+        event_function = get_mock_event_function(
+            mock_smlight_client, SmEvents.FW_UPD_done
+        )
+
+        event_function(MOCK_FIRMWARE_DONE)
+
+        freezer.tick(timedelta(seconds=5))
+        async_fire_time_changed(hass)
+        await hass.async_block_till_done()
+
+        mock_warning.assert_called_once()
+
+
 async def test_update_release_notes(
     hass: HomeAssistant,
     freezer: FrozenDateTimeFactory,