fix: ha discovery workaround race

src/powersensor_local/zeroconf_devices.py

@@ -222,8 +222,11 @@ async def _on_zc_remove(self, mac: str) -> None:
         # Sync to async bridge for callbacks, per asyncio docs
         # ------------------------------------------------------------------
 
-        def _internal_callback(self, coro: _InternalCallback) -> None:
-            """Helper to prevent gc collection of short-lived callback tasks."""
+        def _internal_callback(self, coro: _InternalCallback) -> Task[None]:
+            """Helper to prevent gc collection of short-lived callback tasks.
+
+            Returns the task, so callers can cancel it early if needed.
+            """
             task = asyncio.create_task(coro)
 
             def cleanup(task: Task[None]) -> None:
@@ -238,6 +241,7 @@ def cleanup(task: Task[None]) -> None:
 
             self._internal_callbacks.add(task)
             task.add_done_callback(cleanup)
+            return task
 
 
     class _Listener(_zc.ServiceListener):
@@ -252,6 +256,13 @@ class _Listener(_zc.ServiceListener):
         def __init__(self, owner: PowersensorZeroconfDevices) -> None:
             self._owner = owner
             self._name_to_mac: dict[str, str] = {}
+            # Lookups still waiting for a plug to answer, by service name, so a
+            # removal arriving first can cancel them.
+            self._lookups: dict[str, Task[None]] = {}
+
+        def _forget_lookup(self, name: str, task: Task[None] | None) -> None:
+            if self._lookups.get(name) is task:
+                del self._lookups[name]
 
         def _parse(self, info: AsyncServiceInfo, name: str) -> tuple[str, str, int] | None:
             """Return (mac, ip, port) from a resolved service record, or None."""
@@ -316,8 +327,16 @@ def _resolve(
                     self._owner._internal_callback(on_resolved(*result))
                 return
 
+            if name in self._lookups:
+                # A lookup for this plug is already waiting; it will dispatch.
+                return
+
             async def request() -> None:
-                if not await info.async_request(zc, _REQUEST_TIMEOUT_MS):
+                answered = await info.async_request(zc, _REQUEST_TIMEOUT_MS)
+                # The plug has answered (or not): from here on a removal goes
+                # through the normal debounced path instead of cancelling this.
+                self._forget_lookup(name, asyncio.current_task())
+                if not answered:
                     self._owner._maybe_log(
                         _LogLevel.WARNING,
                         "No response from %s — device will appear on next mDNS announcement",
@@ -328,7 +347,9 @@ async def request() -> None:
                     await on_resolved(*result)
 
             self._owner._maybe_log(_LogLevel.DEBUG, "No cache entry for %s, requesting it", name)
-            self._owner._internal_callback(request())
+            lookup = self._owner._internal_callback(request())
+            self._lookups[name] = lookup
+            lookup.add_done_callback(lambda task: self._forget_lookup(name, task))
 
         def add_service(self, zc: _zc.Zeroconf, type_: str, name: str) -> None:
             self._resolve(zc, type_, name, self._owner._on_zc_add)
@@ -337,14 +358,21 @@ def update_service(self, zc: _zc.Zeroconf, type_: str, name: str) -> None:
             self._resolve(zc, type_, name, self._owner._on_zc_update)
 
         def remove_service(self, zc: _zc.Zeroconf, type_: str, name: str) -> None:
+            lookup = self._lookups.pop(name, None)
+            if lookup is not None:
+                # Removed before it answered, so this lookup never dispatched.
+                # An earlier add may have, which the MAC check below handles.
+                lookup.cancel()
+                self._owner._maybe_log(_LogLevel.DEBUG, "%s removed before it answered; lookup cancelled", name)
+
             mac = self._name_to_mac.pop(name, None)
-            if mac is None:
+            if mac is not None:
+                self._owner._internal_callback(self._owner._on_zc_remove(mac))
+            elif lookup is None:
                 self._owner._maybe_log(
                     _LogLevel.WARNING,
                     "remove_service for %s: MAC not in cache — removal ignored", name,
                 )
-                return
-            self._owner._internal_callback(self._owner._on_zc_remove(mac))
 
 except ImportError as exc:
     _zeroconf_import_error = exc

tests/test_zeroconf_devices.py

@@ -24,10 +24,14 @@
 class FakeServiceInfo(AsyncServiceInfo):
   """An AsyncServiceInfo with a controllable cache and plug response.
 
-  Keyword arguments override the plug's record, as passed to ServiceInfo.
+  If gate is given, the plug only answers once the gate is set.  Keyword
+  arguments override the plug's record, as passed to ServiceInfo.
   """
 
-  def __init__(self, cached: bool, answers: bool, **record: Any) -> None:
+  def __init__(
+    self, cached: bool, answers: bool, gate: asyncio.Event | None = None,
+    **record: Any,
+  ) -> None:
     super().__init__(TYPE, NAME, **{
       'port': 49476,
       'properties': {b'id': b'aabbccddeeff'},
@@ -36,7 +40,9 @@ def __init__(self, cached: bool, answers: bool, **record: Any) -> None:
     })
     self.cached = cached
     self.answers = answers
+    self.gate = gate
     self.requested = False
+    self.requests = 0
 
   def load_from_cache(self, zc: zeroconf.Zeroconf, now: float | None = None) -> bool:
     return self.cached
@@ -45,6 +51,9 @@ async def async_request(
     self, zc: zeroconf.Zeroconf, timeout: float, *args: Any, **kwargs: Any,
   ) -> bool:
     self.requested = True
+    self.requests += 1
+    if self.gate is not None:
+      await self.gate.wait()
     return self.answers
 
 
@@ -54,7 +63,13 @@ async def settle(devices: PowersensorZeroconfDevices) -> None:
     await asyncio.sleep(0)
     if not devices._internal_callbacks:
       return
-    await asyncio.gather(*list(devices._internal_callbacks))
+    await asyncio.gather(*list(devices._internal_callbacks), return_exceptions=True)
+
+
+async def let_run() -> None:
+  """Let scheduled tasks run up to their next wait, without waiting for them."""
+  for _ in range(5):
+    await asyncio.sleep(0)
 
 
 @pytest.mark.asyncio
@@ -105,6 +120,95 @@ async def test_unanswered_request_is_logged(caplog: pytest.LogCaptureFixture) ->
   assert f'No response from {NAME}' in caplog.text
 
 
+@pytest.mark.asyncio
+async def test_removal_before_answer_cancels_lookup(caplog: pytest.LogCaptureFixture) -> None:
+  zc = MagicMock()
+  devices = PowersensorZeroconfDevices(
+    zeroconf_instance=zc, logger=logging.getLogger('test'),
+  )
+  gate = asyncio.Event()
+  info = FakeServiceInfo(cached=False, answers=True, gate=gate)
+  with caplog.at_level(logging.DEBUG), \
+       patch.object(zd, 'AsyncServiceInfo', return_value=info), \
+       patch.object(devices, '_on_zc_add', AsyncMock()) as on_add, \
+       patch.object(devices, '_on_zc_remove', AsyncMock()) as on_remove:
+    listener = zd._Listener(devices)
+    listener.add_service(zc, TYPE, NAME)
+    await let_run()
+    lookup = listener._lookups[NAME]
+    listener.remove_service(zc, TYPE, NAME)
+    gate.set()
+    await settle(devices)
+  assert lookup.cancelled()
+  assert 'removed before it answered' in caplog.text
+  # The plug was never dispatched, so there is nothing to remove either.
+  on_add.assert_not_awaited()
+  on_remove.assert_not_awaited()
+  assert 'removal ignored' not in caplog.text
+
+
+@pytest.mark.asyncio
+async def test_removal_during_update_lookup_removes_known_plug() -> None:
+  zc = MagicMock()
+  devices = PowersensorZeroconfDevices(zeroconf_instance=zc)
+  gate = asyncio.Event()
+  cached = FakeServiceInfo(cached=True, answers=False)
+  uncached = FakeServiceInfo(cached=False, answers=True, gate=gate)
+  with patch.object(zd, 'AsyncServiceInfo', side_effect=[cached, uncached]), \
+       patch.object(devices, '_on_zc_add', AsyncMock()) as on_add, \
+       patch.object(devices, '_on_zc_update', AsyncMock()) as on_update, \
+       patch.object(devices, '_on_zc_remove', AsyncMock()) as on_remove:
+    listener = zd._Listener(devices)
+    listener.add_service(zc, TYPE, NAME)
+    await settle(devices)
+    listener.update_service(zc, TYPE, NAME)
+    await let_run()
+    listener.remove_service(zc, TYPE, NAME)
+    gate.set()
+    await settle(devices)
+  # The update's lookup is cancelled, but the plug was already dispatched by
+  # the add, so it must still be removed.
+  on_add.assert_awaited_once_with('aabbccddeeff', '192.0.2.10', 49476)
+  on_update.assert_not_awaited()
+  on_remove.assert_awaited_once_with('aabbccddeeff')
+
+
+@pytest.mark.asyncio
+async def test_one_lookup_per_plug_at_a_time() -> None:
+  zc = MagicMock()
+  devices = PowersensorZeroconfDevices(zeroconf_instance=zc)
+  gate = asyncio.Event()
+  info = FakeServiceInfo(cached=False, answers=True, gate=gate)
+  with patch.object(zd, 'AsyncServiceInfo', return_value=info), \
+       patch.object(devices, '_on_zc_add', AsyncMock()) as on_add, \
+       patch.object(devices, '_on_zc_update', AsyncMock()) as on_update:
+    listener = zd._Listener(devices)
+    listener.add_service(zc, TYPE, NAME)
+    await let_run()
+    listener.update_service(zc, TYPE, NAME)
+    gate.set()
+    await settle(devices)
+  assert info.requests == 1
+  on_add.assert_awaited_once_with('aabbccddeeff', '192.0.2.10', 49476)
+  on_update.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_removal_after_answer_is_debounced() -> None:
+  zc = MagicMock()
+  devices = PowersensorZeroconfDevices(zeroconf_instance=zc)
+  info = FakeServiceInfo(cached=False, answers=True)
+  with patch.object(zd, 'AsyncServiceInfo', return_value=info), \
+       patch.object(devices, '_on_zc_add', AsyncMock()), \
+       patch.object(devices, '_on_zc_remove', AsyncMock()) as on_remove:
+    listener = zd._Listener(devices)
+    listener.add_service(zc, TYPE, NAME)
+    await settle(devices)
+    listener.remove_service(zc, TYPE, NAME)
+    await settle(devices)
+  on_remove.assert_awaited_once_with('aabbccddeeff')
+
+
 @pytest.mark.asyncio
 @pytest.mark.parametrize('cached', [True, False], ids=['cached', 'requested'])
 @pytest.mark.parametrize(