Refactor SwitchBot Cloud make_device_data (home-assistant#135698)

bachya · Jan 19, 2025 · 41fe863 · 41fe863
1 parent dfc4cdf
commit 41fe863
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 58 deletions.
diff --git a/homeassistant/components/switchbot_cloud/__init__.py b/homeassistant/components/switchbot_cloud/__init__.py
@@ -8,7 +8,7 @@
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_API_KEY, CONF_API_TOKEN, Platform
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import ConfigEntryNotReady
 
 from .const import DOMAIN
@@ -43,77 +43,97 @@ class SwitchbotCloudData:
     devices: SwitchbotDevices
 
 
-@callback
-def prepare_device(
+async def coordinator_for_device(
     hass: HomeAssistant,
     api: SwitchBotAPI,
     device: Device | Remote,
     coordinators_by_id: dict[str, SwitchBotCoordinator],
-) -> tuple[Device | Remote, SwitchBotCoordinator]:
+) -> SwitchBotCoordinator:
     """Instantiate coordinator and adds to list for gathering."""
     coordinator = coordinators_by_id.setdefault(
         device.device_id, SwitchBotCoordinator(hass, api, device)
     )
-    return (device, coordinator)
 
+    if coordinator.data is None:
+        await coordinator.async_config_entry_first_refresh()
 
-@callback
-def make_device_data(
+    return coordinator
+
+
+async def make_switchbot_devices(
     hass: HomeAssistant,
     api: SwitchBotAPI,
     devices: list[Device | Remote],
     coordinators_by_id: dict[str, SwitchBotCoordinator],
 ) -> SwitchbotDevices:
-    """Make device data."""
+    """Make SwitchBot devices."""
     devices_data = SwitchbotDevices()
-    for device in devices:
-        if isinstance(device, Remote) and device.device_type.endswith(
-            "Air Conditioner"
-        ):
-            devices_data.climates.append(
-                prepare_device(hass, api, device, coordinators_by_id)
-            )
-        if (
-            isinstance(device, Device)
-            and (
-                device.device_type.startswith("Plug")
-                or device.device_type in ["Relay Switch 1PM", "Relay Switch 1"]
-            )
-        ) or isinstance(device, Remote):
-            devices_data.switches.append(
-                prepare_device(hass, api, device, coordinators_by_id)
-            )
-        if isinstance(device, Device) and device.device_type in [
-            "Meter",
-            "MeterPlus",
-            "WoIOSensor",
-            "Hub 2",
-            "MeterPro",
-            "MeterPro(CO2)",
-            "Relay Switch 1PM",
-            "Plug Mini (US)",
-            "Plug Mini (JP)",
-        ]:
-            devices_data.sensors.append(
-                prepare_device(hass, api, device, coordinators_by_id)
-            )
-        if isinstance(device, Device) and device.device_type in [
-            "K10+",
-            "K10+ Pro",
-            "Robot Vacuum Cleaner S1",
-            "Robot Vacuum Cleaner S1 Plus",
-        ]:
-            devices_data.vacuums.append(
-                prepare_device(hass, api, device, coordinators_by_id)
-            )
-
-        if isinstance(device, Device) and device.device_type.startswith("Smart Lock"):
-            devices_data.locks.append(
-                prepare_device(hass, api, device, coordinators_by_id)
-            )
+    await gather(
+        *[
+            make_device_data(hass, api, device, devices_data, coordinators_by_id)
+            for device in devices
+        ]
+    )
+
     return devices_data
 
 
+async def make_device_data(
+    hass: HomeAssistant,
+    api: SwitchBotAPI,
+    device: Device | Remote,
+    devices_data: SwitchbotDevices,
+    coordinators_by_id: dict[str, SwitchBotCoordinator],
+) -> None:
+    """Make device data."""
+    if isinstance(device, Remote) and device.device_type.endswith("Air Conditioner"):
+        coordinator = await coordinator_for_device(
+            hass, api, device, coordinators_by_id
+        )
+        devices_data.climates.append((device, coordinator))
+    if (
+        isinstance(device, Device)
+        and (
+            device.device_type.startswith("Plug")
+            or device.device_type in ["Relay Switch 1PM", "Relay Switch 1"]
+        )
+    ) or isinstance(device, Remote):
+        coordinator = await coordinator_for_device(
+            hass, api, device, coordinators_by_id
+        )
+        devices_data.switches.append((device, coordinator))
+
+    if isinstance(device, Device) and device.device_type in [
+        "Meter",
+        "MeterPlus",
+        "WoIOSensor",
+        "Hub 2",
+        "MeterPro",
+        "MeterPro(CO2)",
+        "Relay Switch 1PM",
+        "Plug Mini (US)",
+        "Plug Mini (JP)",
+    ]:
+        devices_data.sensors.append((device, coordinator))
+
+    if isinstance(device, Device) and device.device_type in [
+        "K10+",
+        "K10+ Pro",
+        "Robot Vacuum Cleaner S1",
+        "Robot Vacuum Cleaner S1 Plus",
+    ]:
+        coordinator = await coordinator_for_device(
+            hass, api, device, coordinators_by_id
+        )
+        devices_data.vacuums.append((device, coordinator))
+
+    if isinstance(device, Device) and device.device_type.startswith("Smart Lock"):
+        coordinator = await coordinator_for_device(
+            hass, api, device, coordinators_by_id
+        )
+        devices_data.locks.append((device, coordinator))
+
+
 async def async_setup_entry(hass: HomeAssistant, config: ConfigEntry) -> bool:
     """Set up SwitchBot via API from a config entry."""
     token = config.data[CONF_API_TOKEN]
@@ -131,12 +151,13 @@ async def async_setup_entry(hass: HomeAssistant, config: ConfigEntry) -> bool:
         raise ConfigEntryNotReady from ex
     _LOGGER.debug("Devices: %s", devices)
     coordinators_by_id: dict[str, SwitchBotCoordinator] = {}
+
+    switchbot_devices = await make_switchbot_devices(
+        hass, api, devices, coordinators_by_id
+    )
     hass.data.setdefault(DOMAIN, {})
     hass.data[DOMAIN][config.entry_id] = SwitchbotCloudData(
-        api=api, devices=make_device_data(hass, api, devices, coordinators_by_id)
-    )
-    await gather(
-        *[coordinator.async_refresh() for coordinator in coordinators_by_id.values()]
+        api=api, devices=switchbot_devices
     )
     await hass.config_entries.async_forward_entry_setups(config, PLATFORMS)
     return True

diff --git a/tests/components/switchbot_cloud/test_init.py b/tests/components/switchbot_cloud/test_init.py
@@ -116,7 +116,7 @@ async def test_setup_entry_fails_when_refreshing(
     mock_get_status.side_effect = CannotConnect
     entry = configure_integration(hass)
     await hass.config_entries.async_setup(entry.entry_id)
-    assert entry.state is ConfigEntryState.LOADED
+    assert entry.state is ConfigEntryState.SETUP_RETRY
 
     hass.bus.async_fire(EVENT_HOMEASSISTANT_START)
     await hass.async_block_till_done()