diff --git a/src/powersensor_local/zeroconf_devices.py b/src/powersensor_local/zeroconf_devices.py index 204772d..304b624 100644 --- a/src/powersensor_local/zeroconf_devices.py +++ b/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 diff --git a/tests/test_zeroconf_devices.py b/tests/test_zeroconf_devices.py index 567b7d6..674765f 100644 --- a/tests/test_zeroconf_devices.py +++ b/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(