Wasp in a Box's Wasp Timeout (#550)

* renamed tests.common to tests.helpers

* linting after renaming common to helpers

* added timer helper and corresponding test

* made the patch async call later block into a fixture

* added wasp timeout + tests

* added the options and description to translation file, add a check to not enable the timer if the timeout is 0

* fixing some logic errors with the timer implementation on wiab and unnecessary call for delayed wiab logic

* increasing voodoo loop

* removed requirement box_sensors and wasp_sensors to be present for initialization allowing wiab to be instantiated even with no sensors without errors. this is common during the startup and reloads.

Author: jseidl

custom_components/magic_areas/binary_sensor/wasp_in_a_box.py

@@ -8,20 +8,24 @@
     BinarySensorEntity,
 )
 from homeassistant.const import STATE_OFF, STATE_ON
-from homeassistant.core import Event, EventStateChangedData, callback
+from homeassistant.core import Event, EventStateChangedData, State, callback
 from homeassistant.helpers.event import async_track_state_change_event
 
 from custom_components.magic_areas.base.entities import MagicEntity
 from custom_components.magic_areas.base.magic import MagicArea
 from custom_components.magic_areas.const import (
     CONF_WASP_IN_A_BOX_DELAY,
     CONF_WASP_IN_A_BOX_WASP_DEVICE_CLASSES,
+    CONF_WASP_IN_A_BOX_WASP_TIMEOUT,
     DEFAULT_WASP_IN_A_BOX_DELAY,
     DEFAULT_WASP_IN_A_BOX_WASP_DEVICE_CLASSES,
+    DEFAULT_WASP_IN_A_BOX_WASP_TIMEOUT,
+    ONE_MINUTE,
     WASP_IN_A_BOX_BOX_DEVICE_CLASSES,
     MagicAreasFeatureInfoWaspInABox,
     MagicAreasFeatures,
 )
+from custom_components.magic_areas.helpers.timer import ReusableTimer
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -34,35 +38,36 @@ class AreaWaspInABoxBinarySensor(MagicEntity, BinarySensorEntity):
     """Wasp In The Box logic tracking sensor for the area."""
 
     feature_info = MagicAreasFeatureInfoWaspInABox()
-    _wasp_sensors: list[str]
-    _box_sensors: list[str]
-    delay: int
-    wasp: bool
 
     def __init__(self, area: MagicArea) -> None:
         """Initialize the area presence binary sensor."""
 
         MagicEntity.__init__(self, area, domain=BINARY_SENSOR_DOMAIN)
         BinarySensorEntity.__init__(self)
 
-        self.delay = self.area.feature_config(MagicAreasFeatures.WASP_IN_A_BOX).get(
-            CONF_WASP_IN_A_BOX_DELAY, DEFAULT_WASP_IN_A_BOX_DELAY
-        )
+        self._delay: int = self.area.feature_config(
+            MagicAreasFeatures.WASP_IN_A_BOX
+        ).get(CONF_WASP_IN_A_BOX_DELAY, DEFAULT_WASP_IN_A_BOX_DELAY)
+
+        self._wasp_timeout: int = self.area.feature_config(
+            MagicAreasFeatures.WASP_IN_A_BOX
+        ).get(CONF_WASP_IN_A_BOX_WASP_TIMEOUT, DEFAULT_WASP_IN_A_BOX_WASP_TIMEOUT)
 
         self._attr_device_class = BinarySensorDeviceClass.PRESENCE
         self._attr_extra_state_attributes = {
             ATTR_BOX: STATE_OFF,
             ATTR_WASP: STATE_OFF,
         }
 
-        self.wasp = False
+        self.wasp: bool = False
+        self._wasp_timer: ReusableTimer | None = None
         self._attr_is_on: bool = False
 
-        self._wasp_sensors = []
-        self._box_sensors = []
+        self._wasp_sensors: list[str] = []
+        self._box_sensors: list[str] = []
 
     async def async_added_to_hass(self) -> None:
-        """Call to add the system to hass."""
+        """Call to add the entity to hass."""
         await super().async_added_to_hass()
 
         # Check entities exist
@@ -80,74 +85,91 @@ async def async_added_to_hass(self) -> None:
                 continue
             self._wasp_sensors.append(dc_entity_id)
 
-        if not self._wasp_sensors:
-            raise RuntimeError(f"{self.area.name}: No valid wasp sensors defined.")
-
         for device_class in WASP_IN_A_BOX_BOX_DEVICE_CLASSES:
             dc_entity_id = f"{BINARY_SENSOR_DOMAIN}.magic_areas_aggregates_{self.area.slug}_aggregate_{device_class}"
             dc_state = self.hass.states.get(dc_entity_id)
             if not dc_state:
                 continue
             self._box_sensors.append(dc_entity_id)
 
-        if not self._box_sensors:
-            raise RuntimeError(f"{self.area.name}: No valid wasp sensors defined.")
+        # Initialize timer if timeout configured
+        if self._wasp_timeout > 0:
 
-        # Add listeners
+            async def forget_wasp(now):
+                self.wasp = False
+                self._attr_extra_state_attributes[ATTR_WASP] = STATE_OFF
+                self._attr_is_on = self.wasp
+                self.schedule_update_ha_state()
 
-        self.async_on_remove(
-            async_track_state_change_event(
-                self.hass, self._wasp_sensors, self._wasp_sensor_state_change
+            self._wasp_timer = ReusableTimer(
+                self.hass, self._wasp_timeout * ONE_MINUTE, forget_wasp
             )
-        )
-        self.async_on_remove(
-            async_track_state_change_event(
-                self.hass, self._box_sensors, self._box_sensor_state_change
+
+        # Add listeners
+        if self._wasp_sensors:
+            self.async_on_remove(
+                async_track_state_change_event(
+                    self.hass, self._wasp_sensors, self._async_wasp_sensor_state_change
+                )
             )
-        )
+        if self._box_sensors:
+            self.async_on_remove(
+                async_track_state_change_event(
+                    self.hass, self._box_sensors, self._async_box_sensor_state_change
+                )
+            )
+
+    async def async_will_remove_from_hass(self) -> None:
+        """Call to remove the entity to hass."""
+        if self._wasp_timer:
+            await self._wasp_timer.async_remove()
+        await super().async_will_remove_from_hass()
 
-    def _wasp_sensor_state_change(self, event: Event[EventStateChangedData]) -> None:
+    @callback
+    async def _async_wasp_sensor_state_change(
+        self, event: Event[EventStateChangedData]
+    ) -> None:
         """Register wasp sensor state change event."""
 
+        new_state: State | None = event.data.get("new_state")
+        old_state: State | None = event.data.get("old_state")
+
         # Ignore state reports that aren't really a state change
-        if not event.data["new_state"] or not event.data["old_state"]:
+        if new_state is None or old_state is None:
             return
-        if event.data["new_state"].state == event.data["old_state"].state:
+        if new_state.state == old_state.state:
             return
 
-        self.wasp_in_a_box(wasp_state=event.data["new_state"].state)
+        self.wasp_in_a_box(wasp_state=new_state.state)
 
-    def _box_sensor_state_change(self, event: Event[EventStateChangedData]) -> None:
+    @callback
+    async def _async_box_sensor_state_change(
+        self, event: Event[EventStateChangedData]
+    ) -> None:
         """Register box sensor state change event."""
 
+        new_state: State | None = event.data.get("new_state")
+        old_state: State | None = event.data.get("old_state")
+
         # Ignore state reports that aren't really a state change
-        if not event.data["new_state"] or not event.data["old_state"]:
+        if new_state is None or old_state is None:
             return
-        if event.data["new_state"].state == event.data["old_state"].state:
+        if new_state.state == old_state.state:
             return
 
-        if self.delay:
+        if self._delay:
             self.wasp = False
             self._attr_is_on = self.wasp
-            self._attr_extra_state_attributes[ATTR_BOX] = event.data["new_state"].state
+            self._attr_extra_state_attributes[ATTR_BOX] = new_state.state
             self._attr_extra_state_attributes[ATTR_WASP] = STATE_OFF
             self.schedule_update_ha_state()
-            self.hass.loop.call_soon_threadsafe(
-                self.wasp_in_a_box_delayed,
-                None,
-                event.data["new_state"].state,
+            if self._wasp_timer:
+                self._wasp_timer.cancel()
+            self.hass.loop.call_later(
+                self._delay, self.wasp_in_a_box, None, new_state.state
             )
         else:
-            self.wasp_in_a_box(box_state=event.data["new_state"].state)
-
-    @callback
-    def wasp_in_a_box_delayed(
-        self,
-        wasp_state: str | None = None,
-        box_state: str | None = None,
-    ) -> None:
-        """Call Wasp In A Box Logic function after a delay."""
-        self.hass.loop.call_later(self.delay, self.wasp_in_a_box, wasp_state, box_state)
+            self.wasp_in_a_box(box_state=new_state.state)
 
     def wasp_in_a_box(
         self,
@@ -181,8 +203,16 @@ def wasp_in_a_box(
         # Main Logic
         if wasp_state == STATE_ON:
             self.wasp = True
+            if self._wasp_timer:
+                self._wasp_timer.cancel()
         elif box_state == STATE_ON:
             self.wasp = False
+            if self._wasp_timer:
+                self._wasp_timer.cancel()
+        else:
+            # Wasp is OFF and Box is OFF → start timer
+            if self._wasp_timer and self.wasp:
+                self._wasp_timer.start()
 
         self._attr_extra_state_attributes[ATTR_BOX] = box_state
         self._attr_extra_state_attributes[ATTR_WASP] = wasp_state

custom_components/magic_areas/config_flow.py

@@ -116,6 +116,7 @@
     CONF_TYPE,
     CONF_WASP_IN_A_BOX_DELAY,
     CONF_WASP_IN_A_BOX_WASP_DEVICE_CLASSES,
+    CONF_WASP_IN_A_BOX_WASP_TIMEOUT,
     CONFIG_FLOW_ENTITY_FILTER_BOOL,
     CONFIG_FLOW_ENTITY_FILTER_EXT,
     CONFIGURABLE_AREA_STATE_MAP,
@@ -1273,6 +1274,9 @@ async def async_step_feature_conf_wasp_in_a_box(self, user_input=None):
             CONF_WASP_IN_A_BOX_DELAY: self._build_selector_number(
                 min_value=0, unit_of_measurement="seconds"
             ),
+            CONF_WASP_IN_A_BOX_WASP_TIMEOUT: self._build_selector_number(
+                min_value=0, unit_of_measurement="minutes"
+            ),
             CONF_WASP_IN_A_BOX_WASP_DEVICE_CLASSES: self._build_selector_select(
                 sorted(WASP_IN_A_BOX_WASP_DEVICE_CLASSES), multiple=True
             ),

custom_components/magic_areas/const.py

@@ -640,6 +640,10 @@ class MetaAreaType(StrEnum):
     "delay",
     90,
 )  # cv.positive_int
+CONF_WASP_IN_A_BOX_WASP_TIMEOUT, DEFAULT_WASP_IN_A_BOX_WASP_TIMEOUT = (
+    "wasp_timeout",
+    0,  # 0 = disabled
+)  # cv.positive_int
 CONF_WASP_IN_A_BOX_WASP_DEVICE_CLASSES, DEFAULT_WASP_IN_A_BOX_WASP_DEVICE_CLASSES = (
     "wasp_device_classes",
     [BinarySensorDeviceClass.MOTION, BinarySensorDeviceClass.OCCUPANCY],
@@ -796,6 +800,9 @@ class MetaAreaType(StrEnum):
         vol.Optional(
             CONF_WASP_IN_A_BOX_DELAY, default=DEFAULT_WASP_IN_A_BOX_DELAY
         ): cv.positive_int,
+        vol.Optional(
+            CONF_WASP_IN_A_BOX_WASP_TIMEOUT, default=DEFAULT_WASP_IN_A_BOX_WASP_TIMEOUT
+        ): cv.positive_int,
         vol.Optional(
             CONF_WASP_IN_A_BOX_WASP_DEVICE_CLASSES,
             default=DEFAULT_WASP_IN_A_BOX_WASP_DEVICE_CLASSES,
@@ -1204,6 +1211,11 @@ class MetaAreaType(StrEnum):
 
 OPTIONS_WASP_IN_A_BOX = [
     (CONF_WASP_IN_A_BOX_DELAY, DEFAULT_WASP_IN_A_BOX_DELAY, cv.positive_int),
+    (
+        CONF_WASP_IN_A_BOX_WASP_TIMEOUT,
+        DEFAULT_WASP_IN_A_BOX_WASP_TIMEOUT,
+        cv.positive_int,
+    ),
     (
         CONF_WASP_IN_A_BOX_WASP_DEVICE_CLASSES,
         DEFAULT_WASP_IN_A_BOX_WASP_DEVICE_CLASSES,

custom_components/magic_areas/helpers/timer.py

@@ -0,0 +1,62 @@
+"""Reusable Timer helper for Magic Areas."""
+
+from collections.abc import Awaitable, Callable
+from datetime import datetime
+import logging
+
+from homeassistant.core import CALLBACK_TYPE, HomeAssistant
+from homeassistant.helpers.event import async_call_later
+
+LOGGER: logging.Logger = logging.getLogger(__name__)
+
+
+class ReusableTimer:
+    """Single active reusable timer with fixed delay and callback."""
+
+    def __init__(
+        self,
+        hass: HomeAssistant,
+        delay: float,
+        callback: Callable[[datetime], Awaitable[None]],
+    ) -> None:
+        """Initialize the timer with a fixed delay and async callback."""
+        self.hass = hass
+        self._delay = delay
+        self._callback = callback
+        self._handle: CALLBACK_TYPE | None = None
+        self._token: int = 0  # protects against race conditions
+
+        LOGGER.debug(
+            "Initialized logger with delay=%d and callback=%s",
+            self._delay,
+            str(self._callback),
+        )
+
+    def start(self) -> None:
+        """(Re)start the timer using the configured delay + callback."""
+        self.cancel()
+        self._token += 1
+        token = self._token
+
+        async def _scheduled(now: datetime) -> None:
+            # Ignore if a newer start() happened after scheduling
+            if token != self._token:
+                LOGGER.debug("Token mismatch. Skipping (%d/%d)", self._token, token)
+                return
+            self._handle = None
+            await self._callback(now)
+            LOGGER.debug("Timer fired.")
+
+        self._handle = async_call_later(self.hass, self._delay, _scheduled)
+        LOGGER.debug("Timer started.")
+
+    def cancel(self) -> None:
+        """Cancel the timer if running."""
+        if self._handle:
+            self._handle()  # async_call_later returns a cancel function
+            self._handle = None
+            LOGGER.debug("Timer cancelled.")
+
+    async def async_remove(self) -> None:
+        """Cleanup when entity/integration is removed."""
+        self.cancel()

custom_components/magic_areas/translations/en.json

@@ -178,10 +178,12 @@
         "description": "If we see a wasp, then there is a wasp in the box. If we close the box while there is a wasp in the box, then the wasp remains in the box. If the box remains closed when there is a wasp in the box, and the wasp stops moving, we assume it's still in the box. If the box is opened, and the wasp is not moving, we assume that the wasp has escaped (there is no wasp in the box). This feature monitors your `door` aggregates as the box and your `motion` (default) aggregates as your wasp. A new `binary_sensor` will be created and automatically tracked for presence.",
         "data": {
           "delay": "Delay",
+          "wasp_timeout": "Wasp Timeout",
           "wasp_device_classes": "Wasp sensor device classes"
         },
         "data_description": {
           "delay": "Time to wait before checking for wasps after a box changing state. Set to 0 to disable.",
+          "wasp_timeout": "If greater than 0, wasp will be 'forgotten' after being inactive for this long.",
           "wasp_device_classes": "If you wish other aggregates to be considered as wasps, please select from the list."
         }
       },

tests/conftest.py

@@ -3,6 +3,7 @@
 from collections.abc import AsyncGenerator
 import logging
 from typing import Any
+from unittest.mock import patch
 
 import pytest
 from pytest_homeassistant_custom_component.common import MockConfigEntry
@@ -22,13 +23,14 @@
     AreaType,
 )
 
-from tests.common import (
+from tests.const import DEFAULT_MOCK_AREA, MOCK_AREAS, MockAreaIds
+from tests.helpers import (
     get_basic_config_entry_data,
+    immediate_call_factory,
     init_integration,
     setup_mock_entities,
     shutdown_integration,
 )
-from tests.const import DEFAULT_MOCK_AREA, MOCK_AREAS, MockAreaIds
 from tests.mocks import MockBinarySensor
 
 _LOGGER = logging.getLogger(__name__)
@@ -45,6 +47,19 @@ async def auto_enable_custom_integrations(
     yield
 
 
+# Timer-related
+
+
+@pytest.fixture
+def patch_async_call_later(hass):
+    """Automatically patch async_call_later for ReusableTimer tests."""
+    with patch(
+        "custom_components.magic_areas.helpers.timer.async_call_later",
+        side_effect=immediate_call_factory(hass),
+    ):
+        yield
+
+
 # Config entries