From 41fe863b72090bdb52b009da21ddb6b53b79a7fb Mon Sep 17 00:00:00 2001 From: Mick Montorier-Aberman Date: Sun, 19 Jan 2025 14:22:21 +0100 Subject: [PATCH] Refactor SwitchBot Cloud make_device_data (#135698) --- .../components/switchbot_cloud/__init__.py | 133 ++++++++++-------- tests/components/switchbot_cloud/test_init.py | 2 +- 2 files changed, 78 insertions(+), 57 deletions(-) diff --git a/homeassistant/components/switchbot_cloud/__init__.py b/homeassistant/components/switchbot_cloud/__init__.py index e7313648e6a..e3cfc2172c7 100644 --- a/homeassistant/components/switchbot_cloud/__init__.py +++ b/homeassistant/components/switchbot_cloud/__init__.py @@ -8,7 +8,7 @@ from switchbot_api import CannotConnect, Device, InvalidAuth, Remote, SwitchBotA 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() + + return coordinator -@callback -def make_device_data( +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) - ) + await gather( + *[ + make_device_data(hass, api, device, devices_data, coordinators_by_id) + for device in devices + ] + ) - if isinstance(device, Device) and device.device_type.startswith("Smart Lock"): - devices_data.locks.append( - prepare_device(hass, api, device, coordinators_by_id) - ) 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 index 43431ae04c0..d5728faf369 100644 --- 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()