Fix wirelesstag KeyError when a tag disappears from a reload

homeassistant/components/wirelesstag/binary_sensor.py

@@ -126,6 +126,5 @@ def updated_state_value(self):
     @callback
     def _on_binary_event_callback(self, new_tag: SensorTag) -> None:
         """Update state from arrived push notification."""
-        self._tag = new_tag
-        self._state = self.updated_state_value()
+        self._apply_tag(new_tag)
         self.async_write_ha_state()

homeassistant/components/wirelesstag/entity.py

@@ -39,6 +39,7 @@ def __init__(self, api: WirelessTagPlatform, tag: SensorTag) -> None:
         self.tag_manager_mac = self._tag.tag_manager_mac
         self._attr_name = self._tag.name
         self._state = None
+        self._tag_missing = False
 
     @property
     def principal_value(self):
@@ -63,19 +64,29 @@ def decorate_value(self, value):
     @override
     def available(self) -> bool:
         """Return True if entity is available."""
-        return self._tag.is_alive
+        return not self._tag_missing and self._tag.is_alive
 
     def update(self) -> None:
         """Update state."""
         if not self.should_poll:
             return
 
         updated_tags = self._api.load_tags()
-        if (updated_tag := updated_tags[self._uuid]) is None:
-            _LOGGER.error('Unable to update tag: "%s"', self.name)
+        if (updated_tag := updated_tags.get(self._uuid)) is None:
+            # The tag is gone from the account, so there is nothing left to
+            # poll. Reported through availability rather than a log message
+            # repeated on every poll.
+            _LOGGER.debug('Unable to update tag: "%s"', self.name)
+            self._tag_missing = True
+            self._state = None
             return
 
-        self._tag = updated_tag
+        self._apply_tag(updated_tag)
+
+    def _apply_tag(self, tag: SensorTag) -> None:
+        """Adopt a freshly received tag and refresh the state from it."""
+        self._tag_missing = False
+        self._tag = tag
         self._state = self.updated_state_value()
 
     @property

homeassistant/components/wirelesstag/sensor.py

@@ -161,6 +161,5 @@ def _sensor(self):
     def _update_tag_info_callback(self, new_tag: SensorTag) -> None:
         """Handle push notification sent by tag manager."""
         _LOGGER.debug("Entity to update state: %s with new tag: %s", self, new_tag)
-        self._tag = new_tag
-        self._state = self.updated_state_value()
+        self._apply_tag(new_tag)
         self.async_write_ha_state()

tests/components/wirelesstag/test_sensor.py

@@ -1,13 +1,21 @@
 """Tests for the Wireless Sensor Tags sensor platform."""
 
+from collections.abc import Awaitable, Callable
 from unittest.mock import MagicMock, patch
 
+from freezegun.api import FrozenDateTimeFactory
 import pytest
 
-from homeassistant.const import PERCENTAGE, UnitOfTemperature
+from homeassistant.components.sensor import SCAN_INTERVAL
+from homeassistant.const import PERCENTAGE, STATE_UNAVAILABLE, UnitOfTemperature
 from homeassistant.core import HomeAssistant, valid_entity_id
 from homeassistant.setup import async_setup_component
 
+from tests.common import async_fire_time_changed
+
+UUID = "00000000-0000-0000-0000-000000000001"
+ENTITY_ID = "sensor.wirelesstag_bedroom_temperature"
+
 CONFIG = {
     "wirelesstag": {"username": "foo@bar.com", "password": "secret"},
     "sensor": {
@@ -20,7 +28,7 @@
 def _mock_tag(name: str) -> MagicMock:
     """Return a mocked wirelesstagpy SensorTag exposing temperature and humidity."""
     tag = MagicMock()
-    tag.uuid = "00000000-0000-0000-0000-000000000001"
+    tag.uuid = UUID
     tag.tag_id = 1
     tag.tag_manager_mac = "ABCDEF012345"
     tag.name = name
@@ -46,14 +54,40 @@ def _sensor(sensor_type: str) -> MagicMock:
     return tag
 
 
+async def _poll(hass: HomeAssistant, freezer: FrozenDateTimeFactory) -> None:
+    """Let the platform poll the tag once."""
+    freezer.tick(SCAN_INTERVAL)
+    async_fire_time_changed(hass)
+    await hass.async_block_till_done()
+
+
+async def _recover_by_poll(
+    hass: HomeAssistant,
+    tag: MagicMock,
+    mock_api: MagicMock,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Make the tag available again through the polling path."""
+    mock_api.load_tags.return_value = {tag.uuid: tag}
+    await _poll(hass, freezer)
+
+
+async def _recover_by_push(
+    hass: HomeAssistant,
+    tag: MagicMock,
+    mock_api: MagicMock,
+    freezer: FrozenDateTimeFactory,
+) -> None:
+    """Make the tag available again through a push notification."""
+    push_callback = mock_api.start_monitoring.call_args[0][0]
+    # The library calls back from its own worker thread.
+    await hass.async_add_executor_job(push_callback, {tag.uuid: tag}, {})
+
+
 @pytest.mark.parametrize(
     ("tag_name", "expected_entity_id"),
     [
-        pytest.param(
-            "Bedroom",
-            "sensor.wirelesstag_bedroom_temperature",
-            id="ascii_name",
-        ),
+        pytest.param("Bedroom", ENTITY_ID, id="ascii_name"),
         pytest.param(
             "Küche",
             "sensor.wirelesstag_kuche_temperature",
@@ -87,3 +121,46 @@ async def test_sensor_entity_id_is_valid(
     assert state is not None
     assert valid_entity_id(state.entity_id)
     assert "sets an invalid entity ID" not in caplog.text
+
+
+@pytest.mark.parametrize(
+    "recover",
+    [
+        pytest.param(_recover_by_poll, id="poll"),
+        pytest.param(_recover_by_push, id="push"),
+    ],
+)
+async def test_update_handles_tag_missing_from_reload(
+    hass: HomeAssistant,
+    freezer: FrozenDateTimeFactory,
+    recover: Callable[
+        [HomeAssistant, MagicMock, MagicMock, FrozenDateTimeFactory], Awaitable[None]
+    ],
+) -> None:
+    """Test an update where the tag is no longer returned is handled gracefully.
+
+    If a reload no longer contains the entity's tag, the update must mark the
+    entity unavailable instead of raising a KeyError or keeping a stale value.
+    The entity recovers once the tag shows up again, by poll or by push.
+    """
+    tag = _mock_tag("Bedroom")
+    with patch("homeassistant.components.wirelesstag.WirelessTags") as mock_api_class:
+        mock_api = mock_api_class.return_value
+        mock_api.load_tags.return_value = {tag.uuid: tag}
+
+        assert await async_setup_component(hass, "wirelesstag", CONFIG)
+        await hass.async_block_till_done()
+        assert await async_setup_component(hass, "sensor", CONFIG)
+        await hass.async_block_till_done()
+        assert hass.states.get(ENTITY_ID).state == "21.5"
+
+        # The tag is no longer returned by a reload.
+        mock_api.load_tags.return_value = {}
+        await _poll(hass, freezer)
+
+        assert hass.states.get(ENTITY_ID).state == STATE_UNAVAILABLE
+
+        await recover(hass, tag, mock_api, freezer)
+        await hass.async_block_till_done()
+
+    assert hass.states.get(ENTITY_ID).state == "21.5"