ucapi-framework 1.9.6__tar.gz → 1.9.8__tar.gz

{ucapi_framework-1.9.6 → ucapi_framework-1.9.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ucapi-framework
-Version: 1.9.6
+Version: 1.9.8
 Summary: ucapi framework that provides core functionality for building integrations.
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{ucapi_framework-1.9.6 → ucapi_framework-1.9.8}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ucapi-framework"
-version = "1.9.6"
+version = "1.9.8"
 description = "ucapi framework that provides core functionality for building integrations."
 readme = "README.md"
 requires-python = ">=3.11"

{ucapi_framework-1.9.6 → ucapi_framework-1.9.8}/tests/test_device.py

@@ -771,6 +771,63 @@ class TestWebSocketDevice:
 
         await device.disconnect()
 
+    @pytest.mark.asyncio
+    async def test_normal_close_cleans_up_before_reconnect(
+        self, mock_device_config, event_loop
+    ):
+        """A clean close must not leak the old socket or reconnect in a tight loop."""
+
+        class ClosingWebSocketDevice(ConcreteWebSocketDevice):
+            def __init__(self, *args, **kwargs):
+                super().__init__(*args, **kwargs)
+                self.attempt_times = []
+                self.closed_count = 0
+                self.second_attempt = asyncio.Event()
+
+            async def create_websocket(self):
+                self.attempt_times.append(event_loop.time())
+                if len(self.attempt_times) == 2:
+                    self.second_attempt.set()
+                return Mock()
+
+            async def receive_message(self):
+                return None
+
+            async def close_websocket(self):
+                self.closed_count += 1
+
+        device = ClosingWebSocketDevice(
+            mock_device_config,
+            loop=event_loop,
+            reconnect_interval=0.05,
+            ping_interval=30,
+        )
+
+        await device.connect()
+        await asyncio.wait_for(device.second_attempt.wait(), timeout=1)
+
+        assert device.closed_count >= 1
+        assert device.attempt_times[1] - device.attempt_times[0] >= 0.04
+        await device.disconnect()
+        assert device._ping_task is None
+        assert device._ws is None
+
+    @pytest.mark.asyncio
+    async def test_single_connection_stops_ping_on_normal_close(
+        self, mock_device_config, event_loop
+    ):
+        """A one-shot connection releases its ping task when the peer closes."""
+        device = ConcreteWebSocketDevice(
+            mock_device_config, loop=event_loop, reconnect=False
+        )
+
+        await device.connect()
+        await asyncio.wait_for(device._ws_task, timeout=1)
+
+        assert device.ws_closed is True
+        assert device._ping_task is None
+        assert device._ws is None
+
     @pytest.mark.asyncio
     async def test_websocket_connection_error_no_reconnect(
         self, mock_device_config, event_loop

{ucapi_framework-1.9.6 → ucapi_framework-1.9.8}/tests/test_driver.py

@@ -6,6 +6,8 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import pytest
 import ucapi
+import ucapi.api as uc
+from ucapi.api_definitions import EventCategory
 from ucapi import media_player
 
 from ucapi_framework.device import BaseDeviceInterface, DeviceEvents
@@ -1469,6 +1471,7 @@ class TestHubBasedIntegration:
         await driver.on_subscribe_entities(["media_player.dev1"])
 
         assert driver.api.available_entities.contains("media_player.dev1")
+        assert driver.api.configured_entities.contains("media_player.dev1")
 
     @pytest.mark.asyncio
     async def test_on_subscribe_entities_hub_based_no_config(self):
@@ -3548,3 +3551,491 @@ class TestFilterEntitiesByType:
         # Should only get one instance
         assert len(players) == 1
         assert players[0].id == "media_player.dev1"
+
+
+class TestSubscriptionLifecycle:
+    """Exercise subscription ordering with ucapi's real API and entity stores."""
+
+    @pytest.fixture
+    async def subscription_driver(self):
+        driver = ConcreteDriver(
+            DeviceForTests,
+            [media_player.MediaPlayer],
+            require_connection_before_registry=True,
+            loop=asyncio.get_running_loop(),
+        )
+        driver.api.set_device_state = AsyncMock()
+        driver.api._broadcast_ws_event = AsyncMock()
+        driver.api.configured_entities.add_listener(
+            ucapi.Events.ENTITY_ATTRIBUTES_UPDATED,
+            driver.api._on_entity_attributes_updated,
+        )
+        driver.config_manager = MagicMock()
+        driver.config_manager.get.return_value = DeviceConfigForTests(
+            "dev1", "Device 1", "192.168.1.1"
+        )
+        return driver
+
+    async def drain(self, driver):
+        await driver.api._events.wait_for_complete()
+        for device in list(driver._device_instances.values()):
+            await device.events.wait_for_complete()
+        tasks = getattr(driver, "_subscription_refresh_tasks", {})
+        if tasks:
+            await asyncio.gather(*list(tasks.values()))
+        await driver.api.configured_entities._events.wait_for_complete()
+
+    @pytest.mark.parametrize("already_connected", [False, True])
+    async def test_subscribe_before_registration(
+        self, subscription_driver, already_connected
+    ):
+        driver = subscription_driver
+        config = driver.config_manager.get.return_value
+        device = driver._add_device_instance(config)
+        device.state = "on"
+        if already_connected:
+            await device.connect()
+            await device.events.wait_for_complete()
+
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+
+        entity = driver.api.configured_entities.get("media_player.dev1")
+        assert entity is driver.api.available_entities.get("media_player.dev1")
+        assert entity is not None
+        assert (
+            entity.attributes[media_player.Attributes.STATE] == media_player.States.ON
+        )
+        driver.api._broadcast_ws_event.assert_any_await(
+            uc.WsMsgEvents.ENTITY_CHANGE,
+            {
+                "entity_id": entity.id,
+                "entity_type": entity.entity_type,
+                "attributes": {media_player.Attributes.STATE: media_player.States.ON},
+            },
+            EventCategory.ENTITY,
+        )
+
+        entity.command = AsyncMock(return_value=ucapi.StatusCodes.OK)
+        driver.api.acknowledge_command = AsyncMock()
+        await driver.api._entity_command(
+            None, 1, {"entity_id": entity.id, "cmd_id": media_player.Commands.ON}
+        )
+        entity.command.assert_awaited_once()
+        driver.api.acknowledge_command.assert_awaited_once_with(
+            None, 1, ucapi.StatusCodes.OK
+        )
+
+    async def test_dynamic_discovery_completes_only_requested_entities(
+        self, subscription_driver
+    ):
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+
+        requested = media_player.MediaPlayer("media_player.dev1", "Requested", [], {})
+        other = media_player.MediaPlayer("media_player.dev1.other", "Other", [], {})
+        driver.add_entities([requested, other])
+        await self.drain(driver)
+
+        assert driver.api.configured_entities.get(requested.id) is requested
+        assert not driver.api.configured_entities.contains(other.id)
+
+    async def test_unsubscribe_before_discovery(self, subscription_driver):
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        client = object()
+        await driver.api._subscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        await driver.api._unsubscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+
+        entity = media_player.MediaPlayer("media_player.dev1", "Late", [], {})
+        driver.add_entity(entity)
+        await self.drain(driver)
+        assert not driver.api.configured_entities.contains(entity.id)
+
+    async def test_closed_client_cannot_complete_pending_subscription(
+        self, subscription_driver
+    ):
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        client = object()
+        await driver.api._subscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        driver.api._events.emit(ucapi.Events.CLIENT_DISCONNECTED, websocket=client)
+        await self.drain(driver)
+
+        entity = media_player.MediaPlayer("media_player.dev1", "Late", [], {})
+        driver.add_entity(entity)
+        await self.drain(driver)
+        assert not driver.api.configured_entities.contains(entity.id)
+
+    async def test_subscription_with_lazily_created_standard_device(
+        self, subscription_driver
+    ):
+        driver = subscription_driver
+        driver._require_connection_before_registry = False
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        assert driver.api.configured_entities.contains("media_player.dev1")
+
+    async def test_unsubscribe_before_callback_starts(self, subscription_driver):
+        driver = subscription_driver
+        client = object()
+        driver.on_subscribe_entities = AsyncMock(wraps=driver.on_subscribe_entities)
+        await driver.api._subscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await driver.api._unsubscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        driver.on_subscribe_entities.assert_not_awaited()
+        assert not driver._device_instances
+
+    @pytest.mark.parametrize(
+        "withdraw", ["unsubscribe", "disconnect", "remove", "clear"]
+    )
+    async def test_withdraw_during_connection(self, subscription_driver, withdraw):
+        driver = subscription_driver
+        device = driver._add_device_instance(driver.config_manager.get.return_value)
+        started, release = asyncio.Event(), asyncio.Event()
+        connect = device.connect
+
+        async def delayed_connect():
+            started.set()
+            await release.wait()
+            return await connect()
+
+        device.connect = delayed_connect
+        client = object()
+        await driver.api._subscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await asyncio.wait_for(started.wait(), timeout=1)
+        if withdraw == "unsubscribe":
+            disconnected = asyncio.Event()
+            disconnect = device.disconnect
+
+            async def disconnect_and_signal():
+                await disconnect()
+                disconnected.set()
+
+            device.disconnect = disconnect_and_signal
+            await driver.api._unsubscribe_events(
+                client, {"entity_ids": ["media_player.dev1"]}
+            )
+            await asyncio.wait_for(disconnected.wait(), timeout=1)
+        elif withdraw == "disconnect":
+            driver.api._events.emit(ucapi.Events.CLIENT_DISCONNECTED, websocket=client)
+        elif withdraw == "remove":
+            driver.remove_device("dev1")
+        else:
+            driver.clear_devices()
+        release.set()
+        await self.drain(driver)
+        await device.events.wait_for_complete()
+
+        assert not driver.api.configured_entities.contains("media_player.dev1")
+        assert "media_player.dev1" not in driver._entity_subscriptions
+        if withdraw != "disconnect":
+            assert not device.is_connected
+
+    async def test_background_connection_registers_when_ready(
+        self, subscription_driver
+    ):
+        driver = subscription_driver
+        device = driver._add_device_instance(driver.config_manager.get.return_value)
+        connect = device.connect
+        device.connect = AsyncMock(return_value=True)
+        driver.create_entities = Mock(wraps=driver.create_entities)
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        driver.create_entities.assert_not_called()
+
+        device.state = "on"
+        await connect()
+        await self.drain(driver)
+        assert driver.api.configured_entities.contains("media_player.dev1")
+        driver.create_entities.assert_called_once()
+
+    async def test_failed_connection_can_complete_after_recovery(
+        self, subscription_driver
+    ):
+        driver = subscription_driver
+        device = driver._add_device_instance(driver.config_manager.get.return_value)
+        driver._ensure_device_connected = AsyncMock(return_value=False)
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        assert not driver.api.configured_entities.contains("media_player.dev1")
+
+        await device.connect()
+        await self.drain(driver)
+        assert driver.api.configured_entities.contains("media_player.dev1")
+
+    async def test_missing_sibling_is_registered(self, subscription_driver):
+        driver = subscription_driver
+        device = driver._add_device_instance(driver.config_manager.get.return_value)
+        await device.connect()
+        await device.events.wait_for_complete()
+        first = media_player.MediaPlayer("media_player.dev1", "First", [], {})
+        second = media_player.MediaPlayer("media_player.dev1.zone2", "Second", [], {})
+        driver.add_entity(first)
+        driver.create_entities = Mock(return_value=[first, second])
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": [first.id, second.id]}
+        )
+        await self.drain(driver)
+        assert driver.api.configured_entities.get(first.id) is first
+        assert driver.api.configured_entities.get(second.id) is second
+
+    async def test_unsubscribe_keeps_device_with_pending_sibling(
+        self, subscription_driver
+    ):
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        client = object()
+        await driver.api._subscribe_events(
+            client, {"entity_ids": ["media_player.dev1", "media_player.dev1.zone2"]}
+        )
+        await self.drain(driver)
+        await driver.api._unsubscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        assert driver._device_instances["dev1"].is_connected
+
+        second = media_player.MediaPlayer("media_player.dev1.zone2", "Second", [], {})
+        driver.add_entity(second)
+        await self.drain(driver)
+        assert driver.api.configured_entities.get(second.id) is second
+
+    @pytest.mark.parametrize("withdraw", ["unsubscribe", "disconnect"])
+    async def test_other_client_keeps_pending_request(
+        self, subscription_driver, withdraw
+    ):
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        first, second = object(), object()
+        for client in [first, second]:
+            await driver.api._subscribe_events(
+                client, {"entity_ids": ["media_player.dev1"]}
+            )
+            await self.drain(driver)
+        if withdraw == "unsubscribe":
+            await driver.api._unsubscribe_events(
+                first, {"entity_ids": ["media_player.dev1"]}
+            )
+        else:
+            driver.api._events.emit(ucapi.Events.CLIENT_DISCONNECTED, websocket=first)
+        await self.drain(driver)
+
+        entity = media_player.MediaPlayer("media_player.dev1", "Late", [], {})
+        driver.add_entity(entity)
+        await self.drain(driver)
+        assert driver.api.configured_entities.get(entity.id) is entity
+
+    async def test_resubscribe_after_unsubscribe(self, subscription_driver):
+        driver = subscription_driver
+        client = object()
+        await driver.api._subscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        await driver.api._unsubscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        assert not driver.api.configured_entities.contains("media_player.dev1")
+
+        await driver.api._subscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        assert driver.api.configured_entities.contains("media_player.dev1")
+        assert driver._device_instances["dev1"].is_connected
+
+    async def test_override_can_register_directly_with_ucapi(self, subscription_driver):
+        driver = subscription_driver
+        device = driver._add_device_instance(driver.config_manager.get.return_value)
+        entity = media_player.MediaPlayer("media_player.dev1", "Override", [], {})
+
+        async def subscribe(entity_ids):
+            await device.connect()
+            driver.api.available_entities.add(entity)
+
+        driver.on_subscribe_entities = AsyncMock(side_effect=subscribe)
+        await driver.api._subscribe_events(object(), {"entity_ids": [entity.id]})
+        await self.drain(driver)
+        driver.on_subscribe_entities.assert_awaited_once_with([entity.id])
+        assert driver.api.configured_entities.get(entity.id) is entity
+
+    async def test_cancel_refresh_on_unsubscribe(self, subscription_driver):
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        client = object()
+        await driver.api._subscribe_events(
+            client, {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        started, cancelled = asyncio.Event(), asyncio.Event()
+
+        async def refresh(entity_id):
+            started.set()
+            try:
+                await asyncio.Future()
+            finally:
+                cancelled.set()
+
+        driver.refresh_entity_state = AsyncMock(side_effect=refresh)
+        entity = media_player.MediaPlayer("media_player.dev1", "Late", [], {})
+        driver.add_entity(entity)
+        await asyncio.wait_for(started.wait(), timeout=1)
+        await driver.api._unsubscribe_events(client, {"entity_ids": [entity.id]})
+        await asyncio.wait_for(cancelled.wait(), timeout=1)
+        await self.drain(driver)
+        assert not driver._subscription_refresh_tasks
+
+    async def test_coordinator_refreshes_state_missed_before_discovery(
+        self, subscription_driver
+    ):
+        from ucapi_framework.entity import Entity as FrameworkEntity
+
+        class Player(media_player.MediaPlayer, FrameworkEntity):
+            async def sync_state(self):
+                self.update({media_player.Attributes.VOLUME: self._device.volume})
+
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        device = driver._device_instances["dev1"]
+        device.volume = 42
+        device.push_update()  # No entity existed when this update was emitted.
+        await device.events.wait_for_complete()
+
+        entity = Player("media_player.dev1", "Player", [], {})
+        entity._device = device
+        entity.subscribe_to_device(device)
+        driver.add_entity(entity)
+        await self.drain(driver)
+        assert (
+            driver.api.configured_entities.get(entity.id).attributes[
+                media_player.Attributes.VOLUME
+            ]
+            == 42
+        )
+
+        device.volume = 43
+        device.push_update()
+        await self.drain(driver)
+        assert entity.attributes[media_player.Attributes.VOLUME] == 43
+
+    async def test_subscribe_and_connect_messages_can_overlap(
+        self, subscription_driver
+    ):
+        driver = subscription_driver
+        device = driver._add_device_instance(driver.config_manager.get.return_value)
+        started, both_started, release = (
+            asyncio.Event(),
+            asyncio.Event(),
+            asyncio.Event(),
+        )
+        connect = device.connect
+        calls = 0
+
+        async def delayed_connect():
+            nonlocal calls
+            calls += 1
+            started.set()
+            if calls == 2:
+                both_started.set()
+            await release.wait()
+            return await connect()
+
+        device.connect = delayed_connect
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": ["media_player.dev1"]}
+        )
+        await asyncio.wait_for(started.wait(), timeout=1)
+        await driver.api._handle_ws_event_msg(None, uc.WsMsgEvents.CONNECT, None)
+        await asyncio.wait_for(both_started.wait(), timeout=1)
+        release.set()
+        await self.drain(driver)
+
+        assert len(driver.api.configured_entities.get_all()) == 1
+        assert driver.api.configured_entities.contains("media_player.dev1")
+
+    async def test_shared_subscription_replaces_in_flight_refresh(
+        self, subscription_driver
+    ):
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        first, second = object(), object()
+        for client in [first, second]:
+            await driver.api._subscribe_events(
+                client, {"entity_ids": ["media_player.dev1"]}
+            )
+            await self.drain(driver)
+
+        started, cancelled = asyncio.Event(), asyncio.Event()
+
+        async def refresh(entity_id):
+            started.set()
+            try:
+                await asyncio.Future()
+            finally:
+                cancelled.set()
+
+        driver.refresh_entity_state = AsyncMock(side_effect=refresh)
+        entity = media_player.MediaPlayer("media_player.dev1", "Late", [], {})
+        driver.add_entity(entity)
+        await asyncio.wait_for(started.wait(), timeout=1)
+
+        # ucapi removes the shared configured entry, but the second client still
+        # needs it. The replacement refresh must not leave the old task running.
+        driver.refresh_entity_state = AsyncMock()
+        await driver.api._unsubscribe_events(first, {"entity_ids": [entity.id]})
+        await asyncio.wait_for(cancelled.wait(), timeout=1)
+        await self.drain(driver)
+        assert driver.api.configured_entities.get(entity.id) is entity
+        assert not driver._subscription_refresh_tasks
+
+    async def test_refresh_failure_is_logged_and_task_cleaned_up(
+        self, subscription_driver, caplog
+    ):
+        driver = subscription_driver
+        driver.create_entities = Mock(return_value=[])
+        await driver.api._subscribe_events(
+            object(), {"entity_ids": ["media_player.dev1"]}
+        )
+        await self.drain(driver)
+        driver.refresh_entity_state = AsyncMock(
+            side_effect=RuntimeError("refresh failed")
+        )
+        driver.add_entity(media_player.MediaPlayer("media_player.dev1", "Late", [], {}))
+        tasks = list(driver._subscription_refresh_tasks.values())
+        await asyncio.gather(*tasks, return_exceptions=True)
+        assert not driver._subscription_refresh_tasks
+        assert "Failed to refresh subscribed entity media_player.dev1" in caplog.text
+        assert driver.api.configured_entities.contains("media_player.dev1")

{ucapi_framework-1.9.6 → ucapi_framework-1.9.8}/ucapi_framework/device.py

@@ -572,15 +572,6 @@ class WebSocketDevice(BaseDeviceInterface):
         _LOG.debug("[%s] Disconnecting WebSocket", self.log_id)
         self._stop_ws.set()
 
-        # Stop ping task
-        if self._ping_task and not self._ping_task.done():
-            self._ping_task.cancel()
-            try:
-                await self._ping_task
-            except asyncio.CancelledError:
-                pass
-        self._ping_task = None
-
         # Stop connection task
         if self._ws_task and not self._ws_task.done():
             self._ws_task.cancel()
@@ -589,17 +580,31 @@ class WebSocketDevice(BaseDeviceInterface):
             except asyncio.CancelledError:
                 pass
 
-        # Close WebSocket
+        await self._cleanup_websocket()
+        self._ws_task = None
+        self.events.emit(DeviceEvents.DISCONNECTED, self.identifier)
+
+    async def _cleanup_websocket(self) -> None:
+        """Release the ping task and socket after any connection outcome."""
+        self._is_connected = False
+
+        if self._ping_task:
+            ping_task = self._ping_task
+            self._ping_task = None
+            if not ping_task.done():
+                ping_task.cancel()
+            try:
+                await ping_task
+            except asyncio.CancelledError:
+                pass
+
         if self._ws:
             try:
                 await self.close_websocket()
             except Exception as err:  # pylint: disable=broad-exception-caught
                 _LOG.debug("[%s] Error closing WebSocket: %s", self.log_id, err)
-            self._ws = None
-
-        self._ws_task = None
-        self._is_connected = False
-        self.events.emit(DeviceEvents.DISCONNECTED, self.identifier)
+            finally:
+                self._ws = None
 
     async def _single_connect(self) -> None:
         """Single connection attempt without reconnection."""
@@ -624,13 +629,7 @@ class WebSocketDevice(BaseDeviceInterface):
             _LOG.error("[%s] WebSocket connection error: %s", self.log_id, err)
             self.events.emit(DeviceEvents.ERROR, self.identifier, str(err))
         finally:
-            self._is_connected = False
-            if self._ws:
-                try:
-                    await self.close_websocket()
-                except Exception:  # pylint: disable=broad-exception-caught
-                    pass
-                self._ws = None
+            await self._cleanup_websocket()
 
     async def _connection_loop(self) -> None:
         """
@@ -668,40 +667,25 @@ class WebSocketDevice(BaseDeviceInterface):
             except Exception as err:  # pylint: disable=broad-exception-caught
                 _LOG.warning("[%s] WebSocket connection error: %s", self.log_id, err)
                 self.events.emit(DeviceEvents.ERROR, self.identifier, str(err))
-                self._is_connected = False
+            finally:
+                await self._cleanup_websocket()
 
-                # Clean up
-                if self._ping_task and not self._ping_task.done():
-                    self._ping_task.cancel()
-                    try:
-                        await self._ping_task
-                    except asyncio.CancelledError:
-                        pass
-                    self._ping_task = None
-
-                if self._ws:
-                    try:
-                        await self.close_websocket()
-                    except Exception:  # pylint: disable=broad-exception-caught
-                        pass
-                    self._ws = None
-
-                # Exponential backoff for reconnection
-                if not self._stop_ws.is_set():
-                    _LOG.debug(
-                        "[%s] Reconnecting in %d seconds",
-                        self.log_id,
-                        self._backoff_current,
-                    )
-                    try:
-                        await asyncio.wait_for(
-                            self._stop_ws.wait(), timeout=self._backoff_current
-                        )
-                    except asyncio.TimeoutError:
-                        pass
-                    self._backoff_current = min(
-                        self._backoff_current * 2, self._reconnect_max
+            # A normal close also needs cleanup and a paced reconnect.
+            if not self._stop_ws.is_set():
+                _LOG.debug(
+                    "[%s] Reconnecting in %s seconds",
+                    self.log_id,
+                    self._backoff_current,
+                )
+                try:
+                    await asyncio.wait_for(
+                        self._stop_ws.wait(), timeout=self._backoff_current
                     )
+                except asyncio.TimeoutError:
+                    pass
+                self._backoff_current = min(
+                    self._backoff_current * 2, self._reconnect_max
+                )
 
     async def _message_loop(self) -> None:
         """Main message loop for receiving WebSocket messages."""

{ucapi_framework-1.9.6 → ucapi_framework-1.9.8}/ucapi_framework/driver.py

@@ -9,7 +9,7 @@ Provides common event handlers and device lifecycle management.
 
 import asyncio
 import logging
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from dataclasses import is_dataclass
 from enum import Enum
 from typing import Any, Generic, TypeVar, TypeAlias, cast
@@ -169,6 +169,10 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
             self._entity_classes = entity_classes
 
         self._device_instances: dict[str, DeviceT] = {}
+        # Keep subscription intent until entities are discovered. None identifies
+        # direct calls to on_subscribe_entities() without a WebSocket client.
+        self._entity_subscriptions: dict[str, set[Any]] = {}
+        self._subscription_refresh_tasks: dict[str, asyncio.Task[None]] = {}
         self._config_manager = None  # Set via config_manager property
         self.entity_id_separator = "."  # Default separator for entity IDs
         self._pending_setup_task: asyncio.Task | None = (
@@ -255,11 +259,118 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
         self.api.listens_to(ucapi.Events.DISCONNECT)(self.on_r2_disconnect_cmd)
         self.api.listens_to(ucapi.Events.ENTER_STANDBY)(self.on_r2_enter_standby)
         self.api.listens_to(ucapi.Events.EXIT_STANDBY)(self.on_r2_exit_standby)
-        self.api.listens_to(ucapi.Events.SUBSCRIBE_ENTITIES)(self.on_subscribe_entities)
+        self.api.listens_to(ucapi.Events.SUBSCRIBE_ENTITIES)(
+            self._on_subscribe_entities
+        )
         self.api.listens_to(ucapi.Events.UNSUBSCRIBE_ENTITIES)(
-            self.on_unsubscribe_entities
+            self._on_unsubscribe_entities
+        )
+        self.api.listens_to(ucapi.Events.CLIENT_DISCONNECTED)(
+            self._on_subscription_client_disconnected
         )
 
+    def _on_subscribe_entities(
+        self, entity_ids: list[str], websocket: Any = None
+    ) -> Coroutine[Any, Any, None]:
+        """Record intent synchronously, before async connection work is scheduled."""
+        for entity_id in entity_ids:
+            self._entity_subscriptions.setdefault(entity_id, set()).add(websocket)
+
+        async def subscribe() -> None:
+            # A later unsubscribe or client disconnect may have arrived before
+            # this callback got a chance to run.
+            requested = [
+                entity_id
+                for entity_id in entity_ids
+                if websocket in self._entity_subscriptions.get(entity_id, set())
+            ]
+            if requested:
+                await self.on_subscribe_entities(requested)
+                # Also support overrides which register directly with ucapi.
+                self._complete_pending_subscriptions()
+
+        return subscribe()
+
+    def _on_unsubscribe_entities(
+        self, entity_ids: list[str], websocket: Any = None
+    ) -> Coroutine[Any, Any, None]:
+        """Withdraw intent before an in-flight connection can register entities."""
+        for entity_id in entity_ids:
+            clients = self._entity_subscriptions.get(entity_id)
+            if clients is not None:
+                clients.discard(websocket)
+            if not clients:
+                self._forget_entity_subscription(entity_id)
+            else:
+                # ucapi removes from its shared configured collection before
+                # emitting this event. Retain requests from other live clients.
+                self._complete_entity_subscription(entity_id)
+
+        async def unsubscribe() -> None:
+            withdrawn = [
+                entity_id
+                for entity_id in entity_ids
+                if entity_id not in self._entity_subscriptions
+            ]
+            if withdrawn:
+                await self.on_unsubscribe_entities(withdrawn)
+
+        return unsubscribe()
+
+    def _on_subscription_client_disconnected(self, websocket: Any = None) -> None:
+        """Discard pending intent belonging to a closed client session."""
+        for entity_id, clients in list(self._entity_subscriptions.items()):
+            clients.discard(websocket)
+            if not clients:
+                self._forget_entity_subscription(entity_id)
+
+    def _forget_entity_subscription(self, entity_id: str) -> None:
+        self._entity_subscriptions.pop(entity_id, None)
+        task = self._subscription_refresh_tasks.pop(entity_id, None)
+        if task is not None:
+            task.cancel()
+
+    def _complete_pending_subscriptions(self) -> None:
+        for entity_id in self._entity_subscriptions:
+            self._complete_entity_subscription(entity_id)
+
+    def _complete_entity_subscription(self, entity_id: str) -> None:
+        """Configure an available, requested entity and publish its current state."""
+        if entity_id not in self._entity_subscriptions:
+            return
+        if self.api.configured_entities.contains(entity_id):
+            return
+        entity = self.api.available_entities.get(entity_id)
+        if entity is None:
+            return
+
+        self.api.configured_entities.add(entity)
+        _LOG.debug("Completed subscription for newly available entity %s", entity_id)
+
+        previous_refresh = self._subscription_refresh_tasks.pop(entity_id, None)
+        if previous_refresh is not None:
+            previous_refresh.cancel()
+
+        async def refresh() -> None:
+            if (
+                entity_id in self._entity_subscriptions
+                and self.api.configured_entities.get(entity_id) is entity
+            ):
+                await self.refresh_entity_state(entity_id)
+
+        task = self._loop.create_task(refresh())
+        self._subscription_refresh_tasks[entity_id] = task
+
+        def finished(task: asyncio.Task[None]) -> None:
+            if self._subscription_refresh_tasks.get(entity_id) is task:
+                self._subscription_refresh_tasks.pop(entity_id)
+            if not task.cancelled() and (error := task.exception()) is not None:
+                _LOG.error(
+                    "Failed to refresh subscribed entity %s", entity_id, exc_info=error
+                )
+
+        task.add_done_callback(finished)
+
     # ========================================================================
     # Remote Two Event Handlers (can be overridden)
     # ========================================================================
@@ -342,6 +453,11 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
         if not entity_ids:
             return
 
+        # Direct callers use the same lifecycle as API events. Event dispatch has
+        # already recorded its client, so avoid adding an anonymous owner there.
+        for entity_id in entity_ids:
+            self._entity_subscriptions.setdefault(entity_id, {None})
+
         # Track which devices existed before this subscribe event
         # Used later to skip state refresh for newly added devices
         devices_before_subscribe = set(self._device_instances.keys())
@@ -352,6 +468,8 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
             processed_devices = set()
 
             for entity_id in entity_ids:
+                if entity_id not in self._entity_subscriptions:
+                    continue
                 device_id = self.device_from_entity_id(entity_id)
                 if device_id is None:
                     _LOG.warning(
@@ -391,6 +509,19 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
                         )
                         continue
 
+                    # Connection work can finish after unsubscription/removal.
+                    if self._device_instances.get(device_id) is not device:
+                        await device.disconnect()
+                        continue
+
+                    # Some device classes only start a background connection;
+                    # those will register when DeviceEvents.CONNECTED arrives.
+                    if not device.is_connected or not any(
+                        self.device_from_entity_id(eid) == device_id
+                        for eid in self._entity_subscriptions
+                    ):
+                        continue
+
                     # After successful connection, register entities from the hub (async)
                     await self.async_register_available_entities(
                         self.get_device_config(device_id), device
@@ -398,19 +529,14 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
                 elif (
                     device
                     and device.is_connected
-                    and self.api.configured_entities.get(entity_id) is None
-                ):
-                    # PATCHED: the device was already connected (e.g. by the
-                    # generic CONNECT handler on Remote restart) before this
-                    # SUBSCRIBE_ENTITIES call arrived, so the transition-only
-                    # registration above never ran. Register now so the
-                    # entity doesn't silently stay unconfigured.
-                    _LOG.info(
-                        "Device %s already connected but entity %s not configured "
-                        "- registering now (post-restart fix)",
-                        device_id,
-                        entity_id,
+                    and any(
+                        self.device_from_entity_id(eid) == device_id
+                        and not self.api.available_entities.contains(eid)
+                        for eid in self._entity_subscriptions
                     )
+                ):
+                    # Inspect all requested entities for this hub, including
+                    # missing siblings when the first entity already exists.
                     await self.async_register_available_entities(
                         self.get_device_config(device_id), device
                     )
@@ -421,6 +547,8 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
             processed_devices = set()
 
             for entity_id in entity_ids:
+                if entity_id not in self._entity_subscriptions:
+                    continue
                 eid_device_id = self.device_from_entity_id(entity_id)
                 if eid_device_id is None:
                     continue
@@ -452,10 +580,18 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
                         entity_id,
                     )
 
-        # Refresh entity state only for devices that already existed and are connected
+        self._complete_pending_subscriptions()
+
+        # Newly completed subscriptions have their own initial refresh task.
+        # Refresh other entities only for devices that already existed and are connected
         # Newly added devices will get their state refreshed in on_device_connected()
         # after the device has finished connecting and populated its attributes
         for entity_id in entity_ids:
+            if (
+                entity_id not in self._entity_subscriptions
+                or entity_id in self._subscription_refresh_tasks
+            ):
+                continue
             device_id = self.device_from_entity_id(entity_id)
             if device_id and device_id in devices_before_subscribe:
                 device = self._device_instances.get(device_id)
@@ -479,6 +615,9 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
         """
         _LOG.debug("Unsubscribe entities event: %s", entity_ids)
 
+        for entity_id in entity_ids:
+            self._forget_entity_subscription(entity_id)
+
         # Track which devices need to be checked
         devices_to_check = set()
 
@@ -494,8 +633,12 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
                 self.api.configured_entities.get(entity_id) is not None
                 for entity_id in device_entities
             )
+            any_entity_requested = any(
+                self.device_from_entity_id(entity_id) == device_id
+                for entity_id in self._entity_subscriptions
+            )
 
-            if not any_entity_configured:
+            if not any_entity_configured and not any_entity_requested:
                 # No entities are configured anymore, disconnect and cleanup
                 device = self._device_instances.get(device_id)
                 if device:
@@ -504,6 +647,11 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
                         device_id,
                     )
                     await device.disconnect()
+                    if self._device_instances.get(device_id) is not device or any(
+                        self.device_from_entity_id(entity_id) == device_id
+                        for entity_id in self._entity_subscriptions
+                    ):
+                        continue
                     device.events.remove_all_listeners()
                     self._device_instances.pop(device_id, None)
 
@@ -890,6 +1038,7 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
 
         # Add to available entities
         self.api.available_entities.add(entity)
+        self._complete_entity_subscription(entity.id)
         _LOG.info("Dynamically added entity: %s", entity.id)
 
     def add_entities(
@@ -1156,6 +1305,7 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
             if self.api.available_entities.contains(entity.id):
                 self.api.available_entities.remove(entity.id)
             self.api.available_entities.add(entity)
+            self._complete_entity_subscription(entity.id)
 
     async def async_register_available_entities(
         self, device_config: ConfigT, device: DeviceT
@@ -1247,13 +1397,32 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
         """
         _LOG.debug("Device connected: %s", device_id)
 
-        if device_id not in self._device_instances:
+        device = self._device_instances.get(device_id)
+        if device is None:
             _LOG.warning("Device %s is not configured", device_id)
             return
 
         await self.api.set_device_state(ucapi.DeviceStates.CONNECTED)
 
+        if self._device_instances.get(device_id) is not device:
+            return
+        if (
+            device.is_connected
+            and self._require_connection_before_registry
+            and any(
+                self.device_from_entity_id(entity_id) == device_id
+                and not self.api.available_entities.contains(entity_id)
+                for entity_id in self._entity_subscriptions
+            )
+        ):
+            await self.async_register_available_entities(device.device_config, device)
+            if self._device_instances.get(device_id) is not device:
+                return
+        self._complete_pending_subscriptions()
+
         for entity_id in self.get_entity_ids_for_device(device_id):
+            if entity_id in self._subscription_refresh_tasks:
+                continue
             configured_entity = self.api.configured_entities.get(entity_id)
             if isinstance(configured_entity, FrameworkEntity) and (
                 type(configured_entity).sync_state is not FrameworkEntity.sync_state
@@ -2143,6 +2312,10 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
 
         :param device_id: Device identifier
         """
+        for entity_id in list(self._entity_subscriptions):
+            if self.device_from_entity_id(entity_id) == device_id:
+                self._forget_entity_subscription(entity_id)
+
         if device_id in self._device_instances:
             _LOG.info("Removing device %s", device_id)
             device = self._device_instances.pop(device_id)
@@ -2159,6 +2332,8 @@ class BaseIntegrationDriver(Generic[DeviceT, ConfigT]):
     def clear_devices(self) -> None:
         """Remove all configured devices."""
         _LOG.info("Clearing all configured devices")
+        for entity_id in list(self._entity_subscriptions):
+            self._forget_entity_subscription(entity_id)
         for device in self._device_instances.values():
             device.events.remove_all_listeners()
             self._loop.create_task(device.disconnect())

{ucapi_framework-1.9.6 → ucapi_framework-1.9.8}/ucapi_framework.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ucapi-framework
-Version: 1.9.6
+Version: 1.9.8
 Summary: ucapi framework that provides core functionality for building integrations.
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown