test: Update setup flow test to support new update delay

.gitignore

@@ -82,5 +82,4 @@ cython_debug/
 
 # MkDocs
 site/
-.cache/
-uv.lock
+.cache/

tests/test_setup.py

@@ -603,7 +603,7 @@ async def test_handle_additional_configuration_response_default(self, setup_flow
 
     @pytest.mark.asyncio
     async def test_update_device_mode(self, setup_flow, config_manager):
-        """Test update mode (remove then re-add)."""
+        """Test update mode (remove old entry only on successful completion)."""
         # Add device
         device = DeviceConfigForTests("dev1", "Device 1", "192.168.1.1")
         config_manager.add_or_update(device)
@@ -617,8 +617,9 @@ async def test_update_device_mode(self, setup_flow, config_manager):
         )
         result = await setup_flow.handle_driver_setup(user_response)
 
-        # Device should be removed
-        assert not config_manager.contains("dev1")
+        # Device should still be present (not removed until new config is saved)
+        assert config_manager.contains("dev1")
+        assert setup_flow._selected_config_id == "dev1"
 
         # Should show restore prompt or discovery
         assert isinstance(result, RequestUserInput)

ucapi_framework/driver.py

@@ -171,6 +171,9 @@ def __init__(
         self._device_instances: dict[str, DeviceT] = {}
         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 = (
+            None  # Task handle for entity registration during setup
+        )
         self._setup_event_handlers()
 
     @property
@@ -2125,6 +2128,7 @@ def remove_device(self, device_id: str) -> None:
             _LOG.info("Removing device %s", device_id)
             device = self._device_instances.pop(device_id)
             device.events.remove_all_listeners()
+            self._loop.create_task(device.disconnect())
 
             # Remove all associated entities
             for entity_id in self.get_entity_ids_for_device(device_id):
@@ -2138,6 +2142,7 @@ def clear_devices(self) -> None:
         _LOG.info("Clearing all configured devices")
         for device in self._device_instances.values():
             device.events.remove_all_listeners()
+            self._loop.create_task(device.disconnect())
         self._device_instances.clear()
         self.api.configured_entities.clear()
         self.api.available_entities.clear()
@@ -2162,9 +2167,13 @@ def on_device_added(self, device_config: ConfigT | None) -> None:
         _LOG.debug("Device added: %s", self.get_device_id(device_config))
 
         if self._require_connection_before_registry:
-            # Schedule async device addition as a background task
-            self._loop.create_task(self.async_add_configured_device(device_config))
+            # Schedule async device addition and expose the task so the setup flow
+            # can await it before returning SetupComplete.
+            self._pending_setup_task = self._loop.create_task(
+                self.async_add_configured_device(device_config)
+            )
         else:
+            self._pending_setup_task = None
             self.add_configured_device(device_config, connect=False)
 
     def on_device_removed(self, device_config: ConfigT | None) -> None:

ucapi_framework/setup.py

@@ -116,6 +116,9 @@ def __init__(
         self._setup_step = SetupSteps.INIT
         self._add_mode = False
         self._pending_device_config: ConfigT | None = None  # For multi-screen flows
+        self._selected_config_id: str | None = (
+            None  # ID of config entry to remove on successful update
+        )
         self._pre_discovery_data: dict[
             str, Any
         ] = {}  # Store data from pre-discovery screens
@@ -126,6 +129,18 @@ def __init__(
             None  # Previous version for migration check
         )
 
+    @property
+    def selected_config_id(self) -> str | None:
+        """Get the currently selected config ID (for updates)."""
+        return self._selected_config_id
+
+    @property
+    def selected_config_entry(self) -> ConfigT | None:
+        """Get the currently selected config entry (for updates)."""
+        if self._selected_config_id is None:
+            return None
+        return self.config.get(self._selected_config_id)
+
     @classmethod
     def create_handler(
         cls,
@@ -193,6 +208,7 @@ async def handle_driver_setup(self, msg: SetupDriver) -> SetupAction:
         elif isinstance(msg, AbortDriverSetup):
             _LOG.info("Setup was aborted with code: %s", msg.error)
             self._setup_step = SetupSteps.INIT
+            self._selected_config_id = None
             return SetupError()
         else:
             return SetupError()
@@ -411,9 +427,10 @@ async def _handle_configuration_mode(self, msg: UserDataResponse) -> SetupAction
 
             case "update":
                 choice = msg.input_values["choice"]
-                if not self.config.remove(choice):
+                if not self.config.contains(choice):
                     _LOG.warning("Could not update device: %s", choice)
                     return SetupError(error_type=IntegrationSetupError.OTHER)
+                self._selected_config_id = choice
 
                 self._pre_discovery_data = {}
 
@@ -513,6 +530,32 @@ async def _handle_discovery(self) -> RequestUserInput:
         # No devices found, show manual entry
         return await self._handle_manual_entry()
 
+    async def _await_setup_completion(self) -> None:
+        """
+        Wait for entity registration to complete before returning SetupComplete.
+
+        When require_connection_before_registry=True, on_device_added() fires a
+        background task (async_add_configured_device) that connects and registers
+        entities. This helper awaits that task so the Remote doesn't receive
+        SetupComplete before entities are available.
+
+        Falls back to a 1-second sleep when no task is pending (e.g. when
+        require_connection_before_registry=False).
+        """
+        task = getattr(self.driver, "_pending_setup_task", None)
+        if task is not None:
+            _LOG.debug(
+                "Waiting for device connection and entity registration to complete"
+            )
+            try:
+                await task
+            except Exception as err:  # pylint: disable=broad-except
+                _LOG.warning("Device setup task raised an exception: %s", err)
+            finally:
+                self.driver._pending_setup_task = None
+        else:
+            await asyncio.sleep(1)
+
     async def _finalize_device_setup(
         self, device_config: ConfigT, input_values: dict[str, Any]
     ) -> SetupComplete | SetupError | RequestUserInput:
@@ -542,10 +585,16 @@ async def _finalize_device_setup(
             return additional_screen
 
         # No additional screens, save and complete
+        if self._selected_config_id is not None:
+            _LOG.debug(
+                "Removing old config entry before update: %s", self._selected_config_id
+            )
+            self.config.remove(self._selected_config_id)
+            self._selected_config_id = None
         self.config.add_or_update(self._pending_device_config)
         self._pending_device_config = None
 
-        await asyncio.sleep(1)
+        await self._await_setup_completion()
         _LOG.info("Setup completed for %s", self.get_device_name(device_config))
         return SetupComplete()
 
@@ -661,6 +710,7 @@ async def _handle_additional_configuration_response(
             # If it returns SetupError, cleanup and return it
             if isinstance(result, SetupError):
                 self._pending_device_config = None
+                self._selected_config_id = None
                 return result
 
             # If it returns a device config (ConfigT), replace pending and save
@@ -703,11 +753,18 @@ async def _handle_additional_configuration_response(
             )
 
             # Save the device and complete
+            if self._selected_config_id is not None:
+                _LOG.debug(
+                    "Removing old config entry before update: %s",
+                    self._selected_config_id,
+                )
+                self.config.remove(self._selected_config_id)
+                self._selected_config_id = None
             self.config.add_or_update(self._pending_device_config)
             device_name = self.get_device_name(self._pending_device_config)
             self._pending_device_config = None
 
-            await asyncio.sleep(1)
+            await self._await_setup_completion()
             _LOG.info("Setup completed for %s", device_name)
             return SetupComplete()
 
@@ -723,6 +780,7 @@ async def _handle_additional_configuration_response(
                     repr(self._pending_device_config)[:200],
                 )
             self._pending_device_config = None
+            self._selected_config_id = None
             return SetupError(error_type=IntegrationSetupError.OTHER)
 
     def _has_migration_support(self) -> bool: