Add seperate setup flow to coordinator

custom_components/myskoda/coordinator.py

@@ -8,9 +8,10 @@
 
 from aiohttp import ClientError
 from aiohttp.client_exceptions import ClientResponseError
-from homeassistant.config_entries import ConfigEntry
+from homeassistant.config_entries import ConfigEntry, ConfigEntryState
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers.debounce import Debouncer
+from homeassistant.helpers.start import async_at_started
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
 from myskoda import MySkoda, Vehicle
@@ -34,6 +35,7 @@
 from .const import (
     API_COOLDOWN_IN_SECONDS,
     CONF_POLL_INTERVAL,
+    COORDINATORS,
     DEFAULT_FETCH_INTERVAL_IN_MINUTES,
     DOMAIN,
     MAX_STORED_OPERATIONS,
@@ -119,14 +121,76 @@ def __init__(
         self.update_departure_info = self._debounce(self._update_departure_info)
         self._mqtt_connecting: bool = False
 
+    async def _async_get_minimal_data(self) -> Vehicle:
+        """Internal method that fetches only basic vehicle info."""
+        return await self.myskoda.get_partial_vehicle(self.vin, [])
+
+    async def _async_get_vehicle_data(self) -> Vehicle:
+        """Internal method that fetches vehicle data."""
+        if self.data:
+            if (
+                self.data.vehicle.info.device_platform == "MBB"
+                and self.data.vehicle.info.specification.model == "CitigoE iV"
+            ):
+                _LOGGER.debug(
+                    "Detected Citigo iV, requesting only partial update without health"
+                )
+                vehicle = await self.myskoda.get_partial_vehicle(
+                    self.vin,
+                    [
+                        CapabilityId.AIR_CONDITIONING,
+                        CapabilityId.CHARGING,
+                        CapabilityId.PARKING_POSITION,
+                        CapabilityId.STATE,
+                        CapabilityId.TRIP_STATISTICS,
+                    ],
+                )
+            else:
+                vehicle = await self.myskoda.get_vehicle(self.vin)
+        else:
+            vehicle = await self.myskoda.get_vehicle(self.vin)
+        return vehicle
+
     async def _async_update_data(self) -> State:
         vehicle = None
         user = None
         config = self.data.config if self.data and self.data.config else Config()
         operations = self.operations
 
-        if not self.myskoda.mqtt and not self._mqtt_connecting:
-            self.hass.async_create_task(self._mqtt_connect())
+        if self.config.state == ConfigEntryState.SETUP_IN_PROGRESS:
+            _LOGGER.debug("Performing initial data fetch for vin %s", self.vin)
+            try:
+                user = await self.myskoda.get_user()
+                vehicle = await self._async_get_minimal_data()
+            except ClientResponseError as err:
+                handle_aiohttp_error(
+                    "setup user and vehicle", err, self.hass, self.config
+                )
+                raise UpdateFailed("Failed to retrieve initial data during setup")
+
+            def _async_finish_startup(hass, config, vin) -> None:
+                """Tasks to execute when we have finished starting up."""
+                _LOGGER.debug(
+                    "MySkoda has finished starting up. Scheduling post-start tasks for vin %s.",
+                    vin,
+                )
+                try:
+                    coord = hass.data[DOMAIN][config.entry_id][COORDINATORS][vin]
+                    if not coord.myskoda.mqtt and not coord._mqtt_connecting:
+                        config.async_create_background_task(
+                            self.hass, coord._mqtt_connect(), "mqtt"
+                        )
+                except KeyError:
+                    _LOGGER.debug("Could not connect to MQTT. Waiting for regular poll")
+                    pass
+
+            async_at_started(
+                hass=self.hass,
+                at_start_cb=_async_finish_startup(self.hass, self.config, self.vin),  # pyright: ignore[reportArgumentType]
+            )  # Schedule post-setup tasks
+            return State(vehicle, user, config, operations)
+
+        # Regular update
 
         _LOGGER.debug("Performing scheduled update of all data for vin %s", self.vin)
 
@@ -142,28 +206,7 @@ async def _async_update_data(self) -> State:
 
         # Obtain vehicle data.
         try:
-            if self.data:
-                if (
-                    self.data.vehicle.info.device_platform == "MBB"
-                    and self.data.vehicle.info.specification.model == "CitigoE iV"
-                ):
-                    _LOGGER.debug(
-                        "Detected CitigoE iV, requesting only partial update without health"
-                    )
-                    vehicle = await self.myskoda.get_partial_vehicle(
-                        self.vin,
-                        [
-                            CapabilityId.AIR_CONDITIONING,
-                            CapabilityId.CHARGING,
-                            CapabilityId.PARKING_POSITION,
-                            CapabilityId.STATE,
-                            CapabilityId.TRIP_STATISTICS,
-                        ],
-                    )
-                else:
-                    vehicle = await self.myskoda.get_vehicle(self.vin)
-            else:
-                vehicle = await self.myskoda.get_vehicle(self.vin)
+            vehicle = await self._async_get_vehicle_data()
         except ClientResponseError as err:
             handle_aiohttp_error("vehicle", err, self.hass, self.config)
         except ClientError as err:
@@ -177,8 +220,11 @@ async def _mqtt_connect(self) -> None:
         """Connect to MQTT and handle internals."""
         _LOGGER.debug("Connecting to MQTT.")
         self._mqtt_connecting = True
-        await self.myskoda.enable_mqtt()
-        self.myskoda.subscribe(self._on_mqtt_event)
+        try:
+            await self.myskoda.enable_mqtt()
+            self.myskoda.subscribe(self._on_mqtt_event)
+        except Exception:
+            pass
         self._mqtt_connecting = False
 
     async def _on_mqtt_event(self, event: Event) -> None: