Refactor SwitchBot Cloud make_device_data (#135698)

This commit is contained in:
Mick Montorier-Aberman 2025-01-19 14:22:21 +01:00 committed by GitHub
parent dfc4cdf785
commit 41fe863b72
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 78 additions and 57 deletions

View File

@ -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

View File

@ -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()