Tests

Author: andrew-codechimp

custom_components/label_state/binary_sensor.py

@@ -110,7 +110,6 @@ class LabelStateBinarySensor(BinarySensorEntity):
     _attr_should_poll = False
 
     _state_dict: dict[str, str] = {}
-    _tracked_entities: list[str] = []
 
     def __init__(
         self,
@@ -202,21 +201,19 @@ def _async_entity_registry_modified(
             entity_registry = er.async_get(self.hass)
             entity_entry = entity_registry.async_get(data["entity_id"])
             if entity_entry and self._label in entity_entry.labels:
-                # The entity has a label, ensure we listen to it, if not already
-                if entity_entry.entity_id not in self._tracked_entities:
-                    self.async_on_remove(
-                        async_track_state_change_event(
-                            self.hass,
-                            entity_entry.entity_id,
-                            self._async_state_listener,
-                        )
-                    )
-                    self._tracked_entities.append(entity_entry.entity_id)
-                    LOGGER.debug(
-                        "Found label %s in entity %s",
-                        self._label,
+                # The entity has a label, ensure we listen to it
+                self.async_on_remove(
+                    async_track_state_change_event(
+                        self.hass,
                         entity_entry.entity_id,
+                        self._async_state_listener,
                     )
+                )
+                LOGGER.debug(
+                    "Found label %s in entity %s",
+                    self._label,
+                    entity_entry.entity_id,
+                )
 
             self._calc_state()
             self.async_write_ha_state()

tests/test_binary_sensor.py

@@ -33,19 +33,29 @@
         "state_2",
         "state_3",
         "expected_state",
+        "later_expected_state",
     ),
     [
         (
             "unavailable",
             "on",
             "on",
             "on",
+            "on",
         ),
         (
             "on",
             "on",
             "on",
             "off",
+            "off",
+        ),
+        (
+            "off",
+            "off",
+            "unavailable",
+            "off",
+            "on",
         ),
     ],
 )
@@ -55,6 +65,7 @@ async def test_state_sensor(
     state_2: str,
     state_3: str,
     expected_state: str,
+    later_expected_state: str,
     entity_registry: er.EntityRegistry,
     label_registry: lr.LabelRegistry,
 ) -> None:
@@ -127,7 +138,7 @@ async def test_state_sensor(
     state = hass.states.get("binary_sensor.test_state")
 
     assert state is not None
-    assert state.state == expected_state
+    assert state.state == later_expected_state
 
 
 @pytest.mark.parametrize(