Timestamp relay changes with a read-only state listener

Polling only sees the relay's current state and last change, so two or
more changes between polls (e.g. an on/off round trip) made the lag
estimate unknown and paused learning for about half an hour.

Subscribe to the relay's state changes (read-only) and integrate every
recorded change between polls. Because HA delivers these events one
loop iteration late, the current state stays authoritative and fills in
a change not yet recorded. Unknown periods (non on/off states, missed
transitions, event overflow) still restart the estimate with an
explicit error bound. Without the listener, behavior is as before.

HA skips unload callbacks after unexpected setup errors, so setup now
removes the listener itself on any failure and registers it for unload
only once setup succeeds.

Tested on HA Core 2025.12.0 and 2026.2.3.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

tests/test_coordinator.py

@@ -3,6 +3,7 @@ import importlib.util
 import sys
 import types
 import unittest
+import unittest.mock
 from datetime import datetime, timedelta, timezone
 from math import exp, pi, sin
 from pathlib import Path
@@ -35,9 +36,13 @@ def load_integration():
         'homeassistant', 'homeassistant.const', 'homeassistant.core',
         'homeassistant.config_entries', 'homeassistant.helpers',
         'homeassistant.helpers.storage', 'homeassistant.helpers.update_coordinator',
-        'homeassistant.helpers.entity_registry')}
+        'homeassistant.helpers.entity_registry', 'homeassistant.helpers.event')}
     modules['homeassistant.const'].Platform = types.SimpleNamespace(SENSOR='sensor', BINARY_SENSOR='binary_sensor')
     modules['homeassistant.core'].HomeAssistant = object
+    modules['homeassistant.core'].Event = object
+    modules['homeassistant.core'].callback = lambda func: func
+    modules['homeassistant.helpers.event'].async_track_state_change_event = (
+        lambda hass, entity_ids, action: hass.track_state(entity_ids, action))
     modules['homeassistant.config_entries'].ConfigEntry = object
     modules['homeassistant.helpers.storage'].Store = Store
     modules['homeassistant.helpers.update_coordinator'].DataUpdateCoordinator = CoordinatorBase
@@ -64,6 +69,7 @@ class BoundaryTests(unittest.IsolatedAsyncioTestCase):
             data={}, states=types.SimpleNamespace(get=self.states.get),
             config=types.SimpleNamespace(units=types.SimpleNamespace(temperature_unit='°C')),
             services=types.SimpleNamespace(async_call=Mock(side_effect=AssertionError('No service calls permitted'))),
+            track_state=Mock(return_value=Mock(name='unsub_relay')),
             registry=types.SimpleNamespace(entries=[], async_remove=Mock()),
             config_entries=types.SimpleNamespace(
                 async_forward_entry_setups=AsyncMock(), async_unload_platforms=AsyncMock(return_value=True),
@@ -205,6 +211,10 @@ class BoundaryTests(unittest.IsolatedAsyncioTestCase):
         kept = types.SimpleNamespace(platform='heating_mpc', unique_id='entry_sensor.room_predicted_60min', entity_id='sensor.forecast')
         self.hass.registry.entries = [obsolete, kept]
         self.assertTrue(await integration.async_setup_entry(self.hass, self.entry))
+        self.hass.track_state.assert_called_once_with(['switch.boiler'], unittest.mock.ANY)
+        unsub = self.hass.track_state.return_value
+        self.entry.async_on_unload.assert_any_call(unsub)
+        unsub.assert_not_called()
         self.hass.registry.async_remove.assert_called_once_with('sensor.old_forecast')
         await integration.async_reload_entry(self.hass, self.entry)
         self.hass.config_entries.async_reload.assert_awaited_once_with('entry')
@@ -221,6 +231,7 @@ class BoundaryTests(unittest.IsolatedAsyncioTestCase):
         with self.assertRaises(RuntimeError):
             await integration.async_setup_entry(self.hass, self.entry)
         self.assertNotIn('entry', self.hass.data['heating_mpc'])
+        self.hass.track_state.return_value.assert_called_once_with()
 
     async def test_uninstall_removes_storage(self):
         store = Store()
@@ -425,3 +436,65 @@ class BoundaryTests(unittest.IsolatedAsyncioTestCase):
         self.assertGreater(model.on_count, 100)
         self.assertAlmostEqual(model.theta[1], b, delta=0.1 * b)
         self.assertAlmostEqual(model.theta[0], a, delta=0.25 * a)
+
+    async def test_several_relay_changes_between_polls_are_integrated(self):
+        advance = coordinator.advance_lag
+        await self.sample()
+        lag, uncertainty = self.c._lag['lag'], self.c._lag['uncertainty']
+        start = self.now - timedelta(seconds=300)
+        # The listener saw on at +60 s and off at +200 s; the state reads off again.
+        self.c.record_relay_state(start + timedelta(seconds=60), 'on')
+        self.c.record_relay_state(start + timedelta(seconds=200), 'off')
+        self.switch('off', 100)
+        data = await self.sample()
+        lag = advance(advance(advance(lag, False, 60), True, 140), False, 100)
+        self.assertEqual(data['sample_result'], 'relay_changed')
+        self.assertAlmostEqual(self.c._lag['lag'], lag)
+        self.assertAlmostEqual(self.c._lag['uncertainty'], uncertainty * exp(-300 / 1200))
+        self.assertEqual(len(self.c._relay_events), 0)
+        self.assertEqual((await self.sample())['sample_result'], 'learned')
+
+    async def test_queued_relay_event_is_reconciled_from_current_state(self):
+        advance = coordinator.advance_lag
+        await self.sample()
+        lag = self.c._lag['lag']
+        start = self.now - timedelta(seconds=300)
+        # Only the first change reached the listener before this poll.
+        self.c.record_relay_state(start + timedelta(seconds=60), 'on')
+        self.switch('off', 100)
+        await self.sample()
+        lag = advance(advance(advance(lag, False, 60), True, 140), False, 100)
+        self.assertAlmostEqual(self.c._lag['lag'], lag)
+        # The late event for the change already accounted for is ignored next time.
+        self.c.record_relay_state(start + timedelta(seconds=200), 'off')
+        self.assertEqual((await self.sample())['sample_result'], 'learned')
+
+    async def test_relay_outage_or_dropped_events_restart_lag_estimate(self):
+        for events in ((('unavailable', 60), ('off', 120)), ((None, 0),)):
+            with self.subTest(events=events):
+                self.c._lag = self.c._previous = None
+                await self.sample()
+                start = self.now - timedelta(seconds=300)
+                if events[0][0] is None:
+                    self.c._relay_events_dropped = True
+                for state, offset in events:
+                    if state:
+                        self.c.record_relay_state(start + timedelta(seconds=offset), state)
+                self.switch('off', 180)
+                data = await self.sample()
+                self.assertEqual(data['sample_result'], 'relay_changed')
+                self.assertAlmostEqual(data['heat_lag_uncertainty'], round(exp(-180 / 1200) / 2, 3))
+                self.assertFalse(self.c._relay_events_dropped)
+
+    async def test_relay_listener_records_only_state_changes(self):
+        def event(old, new, when):
+            return types.SimpleNamespace(time_fired=when, data={
+                'old_state': old and types.SimpleNamespace(state=old),
+                'new_state': new and types.SimpleNamespace(state=new, last_changed=when)})
+        self.c.async_relay_changed(event('off', 'on', self.now))
+        self.c.async_relay_changed(event('on', 'on', self.now))  # attribute-only update
+        self.c.async_relay_changed(event('on', None, self.now))  # entity removed
+        self.assertEqual(list(self.c._relay_events), [(self.now, 'on'), (self.now, 'removed')])
+        for _ in range(coordinator.MAX_RELAY_EVENTS):
+            self.c.record_relay_state(self.now, 'on')
+        self.assertTrue(self.c._relay_events_dropped)