Update and fix govee light local (home-assistant#166454)

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

Author: Galorhallen

homeassistant/components/govee_light_local/coordinator.py

@@ -98,7 +98,7 @@ async def set_temperature(self, device: GoveeDevice, temperature: int) -> None:
         """Set light color in kelvin."""
         await device.set_temperature(temperature)
 
-    async def set_scene(self, device: GoveeController, scene: str) -> None:
+    async def set_scene(self, device: GoveeDevice, scene: str) -> None:
         """Set light scene."""
         await device.set_scene(scene)
 

homeassistant/components/govee_light_local/light.py

@@ -80,7 +80,6 @@ def __init__(
 
         super().__init__(coordinator)
         self._device = device
-        device.set_update_callback(self._update_callback)
 
         self._attr_unique_id = device.fingerprint
 
@@ -194,9 +193,20 @@ async def async_turn_off(self, **kwargs: Any) -> None:
         await self.coordinator.turn_off(self._device)
         self.async_write_ha_state()
 
+    async def async_added_to_hass(self) -> None:
+        """Register update callback when entity is added."""
+        await super().async_added_to_hass()
+        self._device.set_update_callback(self._update_callback)
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Unregister update callback when entity is removed."""
+        self._device.set_update_callback(None)
+        await super().async_will_remove_from_hass()
+
     @callback
     def _update_callback(self, device: GoveeDevice) -> None:
-        self.async_write_ha_state()
+        if self.hass:
+            self.async_write_ha_state()
 
     def _save_last_color_state(self) -> None:
         color_mode = self.color_mode

tests/components/govee_light_local/test_light.py

@@ -631,6 +631,68 @@ async def test_scene_restore_temperature(
     assert light.attributes["color_temp_kelvin"] == initial_color
 
 
+async def test_update_callback_registered_and_triggers_state_update(
+    hass: HomeAssistant, mock_govee_api: MagicMock
+) -> None:
+    """Test that update callback is registered and triggers state update."""
+    device = GoveeDevice(
+        controller=mock_govee_api,
+        ip="192.168.1.100",
+        fingerprint="asdawdqwdqwd",
+        sku="H615A",
+        capabilities=DEFAULT_CAPABILITIES,
+    )
+    mock_govee_api.devices = [device]
+
+    entry = MockConfigEntry(domain=DOMAIN)
+    entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert device.update_callback is not None
+
+    light = hass.states.get("light.H615A")
+    assert light is not None
+    assert light.state == "off"
+
+    # Mutate device state and fire callback
+    await device.turn_on()
+    device.update_callback(device)
+    await hass.async_block_till_done()
+
+    light = hass.states.get("light.H615A")
+    assert light is not None
+    assert light.state == "on"
+
+
+async def test_update_callback_cleared_on_remove(
+    hass: HomeAssistant, mock_govee_api: MagicMock
+) -> None:
+    """Test that update callback is cleared when entity is removed."""
+    device = GoveeDevice(
+        controller=mock_govee_api,
+        ip="192.168.1.100",
+        fingerprint="asdawdqwdqwd",
+        sku="H615A",
+        capabilities=DEFAULT_CAPABILITIES,
+    )
+    mock_govee_api.devices = [device]
+
+    entry = MockConfigEntry(domain=DOMAIN)
+    entry.add_to_hass(hass)
+
+    assert await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert device.update_callback is not None
+
+    assert await hass.config_entries.async_remove(entry.entry_id)
+    await hass.async_block_till_done()
+
+    assert device.update_callback is None
+
+
 async def test_scene_none(hass: HomeAssistant, mock_govee_api: MagicMock) -> None:
     """Test turn on 'none' scene."""