Fix yolink device unavailable on startup (#72579)

* fetch device state on startup

* Suggest change

* suggest fix

* fix

* fix

* Fix suggest

* suggest fix

homeassistant/components/yolink/__init__.py

@@ -1,25 +1,28 @@
 """The yolink integration."""
 from __future__ import annotations
 
+import asyncio
 from datetime import timedelta
-import logging
 
+import async_timeout
 from yolink.client import YoLinkClient
+from yolink.device import YoLinkDevice
+from yolink.exception import YoLinkAuthFailError, YoLinkClientError
+from yolink.model import BRDP
 from yolink.mqtt_client import MqttClient
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
 from homeassistant.helpers import aiohttp_client, config_entry_oauth2_flow
 
 from . import api
-from .const import ATTR_CLIENT, ATTR_COORDINATOR, ATTR_MQTT_CLIENT, DOMAIN
+from .const import ATTR_CLIENT, ATTR_COORDINATORS, ATTR_DEVICE, ATTR_MQTT_CLIENT, DOMAIN
 from .coordinator import YoLinkCoordinator
 
 SCAN_INTERVAL = timedelta(minutes=5)
 
-_LOGGER = logging.getLogger(__name__)
 
 PLATFORMS = [Platform.BINARY_SENSOR, Platform.SENSOR, Platform.SIREN, Platform.SWITCH]
 
@@ -41,18 +44,63 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 
     yolink_http_client = YoLinkClient(auth_mgr)
     yolink_mqtt_client = MqttClient(auth_mgr)
-    coordinator = YoLinkCoordinator(hass, yolink_http_client, yolink_mqtt_client)
+
-    await coordinator.init_coordinator()
+    def on_message_callback(message: tuple[str, BRDP]) -> None:
+        data = message[1]
+        device_id = message[0]
+        if data.event is None:
+            return
+        event_param = data.event.split(".")
+        event_type = event_param[len(event_param) - 1]
+        if event_type not in (
+            "Report",
+            "Alert",
+            "StatusChange",
+            "getState",
+        ):
+            return
+        resolved_state = data.data
+        if resolved_state is None:
+            return
+        entry_data = hass.data[DOMAIN].get(entry.entry_id)
+        if entry_data is None:
+            return
+        device_coordinators = entry_data.get(ATTR_COORDINATORS)
+        if device_coordinators is None:
+            return
+        device_coordinator = device_coordinators.get(device_id)
+        if device_coordinator is None:
+            return
+        device_coordinator.async_set_updated_data(resolved_state)
+
     try:
-        await coordinator.async_config_entry_first_refresh()
+        async with async_timeout.timeout(10):
-    except ConfigEntryNotReady as ex:
+            device_response = await yolink_http_client.get_auth_devices()
-        _LOGGER.error("Fetching initial data failed: %s", ex)
+            home_info = await yolink_http_client.get_general_info()
+            await yolink_mqtt_client.init_home_connection(
+                home_info.data["id"], on_message_callback
+            )
+    except YoLinkAuthFailError as yl_auth_err:
+        raise ConfigEntryAuthFailed from yl_auth_err
+    except (YoLinkClientError, asyncio.TimeoutError) as err:
+        raise ConfigEntryNotReady from err
 
     hass.data[DOMAIN][entry.entry_id] = {
         ATTR_CLIENT: yolink_http_client,
         ATTR_MQTT_CLIENT: yolink_mqtt_client,
-        ATTR_COORDINATOR: coordinator,
     }
+    auth_devices = device_response.data[ATTR_DEVICE]
+    device_coordinators = {}
+    for device_info in auth_devices:
+        device = YoLinkDevice(device_info, yolink_http_client)
+        device_coordinator = YoLinkCoordinator(hass, device)
+        try:
+            await device_coordinator.async_config_entry_first_refresh()
+        except ConfigEntryNotReady:
+            # Not failure by fetching device state
+            device_coordinator.data = {}
+        device_coordinators[device.device_id] = device_coordinator
+    hass.data[DOMAIN][entry.entry_id][ATTR_COORDINATORS] = device_coordinators
     hass.config_entries.async_setup_platforms(entry, PLATFORMS)
     return True
 

homeassistant/components/yolink/binary_sensor.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 
 from collections.abc import Callable
 from dataclasses import dataclass
+from typing import Any
 
 from yolink.device import YoLinkDevice
 
@@ -16,7 +17,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
 from .const import (
-    ATTR_COORDINATOR,
+    ATTR_COORDINATORS,
     ATTR_DEVICE_DOOR_SENSOR,
     ATTR_DEVICE_LEAK_SENSOR,
     ATTR_DEVICE_MOTION_SENSOR,
@@ -32,7 +33,7 @@ class YoLinkBinarySensorEntityDescription(BinarySensorEntityDescription):
 
     exists_fn: Callable[[YoLinkDevice], bool] = lambda _: True
     state_key: str = "state"
-    value: Callable[[str], bool | None] = lambda _: None
+    value: Callable[[Any], bool | None] = lambda _: None
 
 
 SENSOR_DEVICE_TYPE = [
@@ -47,14 +48,14 @@ SENSOR_TYPES: tuple[YoLinkBinarySensorEntityDescription, ...] = (
         icon="mdi:door",
         device_class=BinarySensorDeviceClass.DOOR,
         name="State",
-        value=lambda value: value == "open",
+        value=lambda value: value == "open" if value is not None else None,
         exists_fn=lambda device: device.device_type in [ATTR_DEVICE_DOOR_SENSOR],
     ),
     YoLinkBinarySensorEntityDescription(
         key="motion_state",
         device_class=BinarySensorDeviceClass.MOTION,
         name="Motion",
-        value=lambda value: value == "alert",
+        value=lambda value: value == "alert" if value is not None else None,
         exists_fn=lambda device: device.device_type in [ATTR_DEVICE_MOTION_SENSOR],
     ),
     YoLinkBinarySensorEntityDescription(
@@ -62,7 +63,7 @@ SENSOR_TYPES: tuple[YoLinkBinarySensorEntityDescription, ...] = (
         name="Leak",
         icon="mdi:water",
         device_class=BinarySensorDeviceClass.MOISTURE,
-        value=lambda value: value == "alert",
+        value=lambda value: value == "alert" if value is not None else None,
         exists_fn=lambda device: device.device_type in [ATTR_DEVICE_LEAK_SENSOR],
     ),
 )
@@ -74,18 +75,20 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up YoLink Sensor from a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][ATTR_COORDINATOR]
+    device_coordinators = hass.data[DOMAIN][config_entry.entry_id][ATTR_COORDINATORS]
-    sensor_devices = [
+    binary_sensor_device_coordinators = [
-        device
+        device_coordinator
-        for device in coordinator.yl_devices
+        for device_coordinator in device_coordinators.values()
-        if device.device_type in SENSOR_DEVICE_TYPE
+        if device_coordinator.device.device_type in SENSOR_DEVICE_TYPE
     ]
     entities = []
-    for sensor_device in sensor_devices:
+    for binary_sensor_device_coordinator in binary_sensor_device_coordinators:
         for description in SENSOR_TYPES:
-            if description.exists_fn(sensor_device):
+            if description.exists_fn(binary_sensor_device_coordinator.device):
                 entities.append(
-                    YoLinkBinarySensorEntity(coordinator, description, sensor_device)
+                    YoLinkBinarySensorEntity(
+                        binary_sensor_device_coordinator, description
+                    )
                 )
     async_add_entities(entities)
 
@@ -99,18 +102,21 @@ class YoLinkBinarySensorEntity(YoLinkEntity, BinarySensorEntity):
         self,
         coordinator: YoLinkCoordinator,
         description: YoLinkBinarySensorEntityDescription,
-        device: YoLinkDevice,
     ) -> None:
         """Init YoLink Sensor."""
-        super().__init__(coordinator, device)
+        super().__init__(coordinator)
         self.entity_description = description
-        self._attr_unique_id = f"{device.device_id} {self.entity_description.key}"
+        self._attr_unique_id = (
-        self._attr_name = f"{device.device_name} ({self.entity_description.name})"
+            f"{coordinator.device.device_id} {self.entity_description.key}"
+        )
+        self._attr_name = (
+            f"{coordinator.device.device_name} ({self.entity_description.name})"
+        )
 
     @callback
-    def update_entity_state(self, state: dict) -> None:
+    def update_entity_state(self, state: dict[str, Any]) -> None:
         """Update HA Entity State."""
         self._attr_is_on = self.entity_description.value(
-            state[self.entity_description.state_key]
+            state.get(self.entity_description.state_key)
         )
         self.async_write_ha_state()

homeassistant/components/yolink/const.py

@@ -5,7 +5,7 @@ MANUFACTURER = "YoLink"
 HOME_ID = "homeId"
 HOME_SUBSCRIPTION = "home_subscription"
 ATTR_PLATFORM_SENSOR = "sensor"
-ATTR_COORDINATOR = "coordinator"
+ATTR_COORDINATORS = "coordinators"
 ATTR_DEVICE = "devices"
 ATTR_DEVICE_TYPE = "type"
 ATTR_DEVICE_NAME = "name"

homeassistant/components/yolink/coordinator.py

@@ -1,22 +1,18 @@
 """YoLink DataUpdateCoordinator."""
 from __future__ import annotations
 
-import asyncio
 from datetime import timedelta
 import logging
 
 import async_timeout
-from yolink.client import YoLinkClient
 from yolink.device import YoLinkDevice
 from yolink.exception import YoLinkAuthFailError, YoLinkClientError
-from yolink.model import BRDP
-from yolink.mqtt_client import MqttClient
 
 from homeassistant.core import HomeAssistant
-from homeassistant.exceptions import ConfigEntryAuthFailed, ConfigEntryNotReady
+from homeassistant.exceptions import ConfigEntryAuthFailed
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import ATTR_DEVICE, ATTR_DEVICE_STATE, DOMAIN
+from .const import ATTR_DEVICE_STATE, DOMAIN
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -24,9 +20,7 @@ _LOGGER = logging.getLogger(__name__)
 class YoLinkCoordinator(DataUpdateCoordinator[dict]):
     """YoLink DataUpdateCoordinator."""
 
-    def __init__(
+    def __init__(self, hass: HomeAssistant, device: YoLinkDevice) -> None:
-        self, hass: HomeAssistant, yl_client: YoLinkClient, yl_mqtt_client: MqttClient
-    ) -> None:
         """Init YoLink DataUpdateCoordinator.
 
         fetch state every 30 minutes base on yolink device heartbeat interval
@@ -35,75 +29,17 @@ class YoLinkCoordinator(DataUpdateCoordinator[dict]):
         super().__init__(
             hass, _LOGGER, name=DOMAIN, update_interval=timedelta(minutes=30)
         )
-        self._client = yl_client
+        self.device = device
-        self._mqtt_client = yl_mqtt_client
-        self.yl_devices: list[YoLinkDevice] = []
-        self.data = {}
 
-    def on_message_callback(self, message: tuple[str, BRDP]):
+    async def _async_update_data(self) -> dict:
-        """On message callback."""
+        """Fetch device state."""
-        data = message[1]
-        if data.event is None:
-            return
-        event_param = data.event.split(".")
-        event_type = event_param[len(event_param) - 1]
-        if event_type not in (
-            "Report",
-            "Alert",
-            "StatusChange",
-            "getState",
-        ):
-            return
-        resolved_state = data.data
-        if resolved_state is None:
-            return
-        self.data[message[0]] = resolved_state
-        self.async_set_updated_data(self.data)
-
-    async def init_coordinator(self):
-        """Init coordinator."""
         try:
             async with async_timeout.timeout(10):
-                home_info = await self._client.get_general_info()
+                device_state_resp = await self.device.fetch_state_with_api()
-                await self._mqtt_client.init_home_connection(
-                    home_info.data["id"], self.on_message_callback
-                )
-            async with async_timeout.timeout(10):
-                device_response = await self._client.get_auth_devices()
-
-        except YoLinkAuthFailError as yl_auth_err:
-            raise ConfigEntryAuthFailed from yl_auth_err
-
-        except (YoLinkClientError, asyncio.TimeoutError) as err:
-            raise ConfigEntryNotReady from err
-
-        yl_devices: list[YoLinkDevice] = []
-
-        for device_info in device_response.data[ATTR_DEVICE]:
-            yl_devices.append(YoLinkDevice(device_info, self._client))
-
-        self.yl_devices = yl_devices
-
-    async def fetch_device_state(self, device: YoLinkDevice):
-        """Fetch Device State."""
-        try:
-            async with async_timeout.timeout(10):
-                device_state_resp = await device.fetch_state_with_api()
-                if ATTR_DEVICE_STATE in device_state_resp.data:
-                    self.data[device.device_id] = device_state_resp.data[
-                        ATTR_DEVICE_STATE
-                    ]
         except YoLinkAuthFailError as yl_auth_err:
             raise ConfigEntryAuthFailed from yl_auth_err
         except YoLinkClientError as yl_client_err:
-            raise UpdateFailed(
+            raise UpdateFailed from yl_client_err
-                f"Error communicating with API: {yl_client_err}"
+        if ATTR_DEVICE_STATE in device_state_resp.data:
-            ) from yl_client_err
+            return device_state_resp.data[ATTR_DEVICE_STATE]
-
+        return {}
-    async def _async_update_data(self) -> dict:
-        fetch_tasks = []
-        for yl_device in self.yl_devices:
-            fetch_tasks.append(self.fetch_device_state(yl_device))
-        if fetch_tasks:
-            await asyncio.gather(*fetch_tasks)
-        return self.data

homeassistant/components/yolink/entity.py

@@ -3,8 +3,6 @@ from __future__ import annotations
 
 from abc import abstractmethod
 
-from yolink.device import YoLinkDevice
-
 from homeassistant.core import callback
 from homeassistant.helpers.entity import DeviceInfo
 from homeassistant.helpers.update_coordinator import CoordinatorEntity
@@ -19,20 +17,24 @@ class YoLinkEntity(CoordinatorEntity[YoLinkCoordinator]):
     def __init__(
         self,
         coordinator: YoLinkCoordinator,
-        device_info: YoLinkDevice,
     ) -> None:
         """Init YoLink Entity."""
         super().__init__(coordinator)
-        self.device = device_info
 
     @property
     def device_id(self) -> str:
         """Return the device id of the YoLink device."""
-        return self.device.device_id
+        return self.coordinator.device.device_id
+
+    async def async_added_to_hass(self) -> None:
+        """Update state."""
+        await super().async_added_to_hass()
+        return self._handle_coordinator_update()
 
     @callback
     def _handle_coordinator_update(self) -> None:
-        data = self.coordinator.data.get(self.device.device_id)
+        """Update state."""
+        data = self.coordinator.data
         if data is not None:
             self.update_entity_state(data)
 
@@ -40,10 +42,10 @@ class YoLinkEntity(CoordinatorEntity[YoLinkCoordinator]):
     def device_info(self) -> DeviceInfo:
         """Return the device info for HA."""
         return DeviceInfo(
-            identifiers={(DOMAIN, self.device.device_id)},
+            identifiers={(DOMAIN, self.coordinator.device.device_id)},
             manufacturer=MANUFACTURER,
-            model=self.device.device_type,
+            model=self.coordinator.device.device_type,
-            name=self.device.device_name,
+            name=self.coordinator.device.device_name,
         )
 
     @callback

homeassistant/components/yolink/sensor.py

@@ -19,7 +19,7 @@ from homeassistant.helpers.entity_platform import AddEntitiesCallback
 from homeassistant.util import percentage
 
 from .const import (
-    ATTR_COORDINATOR,
+    ATTR_COORDINATORS,
     ATTR_DEVICE_DOOR_SENSOR,
     ATTR_DEVICE_MOTION_SENSOR,
     ATTR_DEVICE_TH_SENSOR,
@@ -54,7 +54,9 @@ SENSOR_TYPES: tuple[YoLinkSensorEntityDescription, ...] = (
         state_class=SensorStateClass.MEASUREMENT,
         value=lambda value: percentage.ordered_list_item_to_percentage(
             [1, 2, 3, 4], value
-        ),
+        )
+        if value is not None
+        else None,
         exists_fn=lambda device: device.device_type
         in [ATTR_DEVICE_DOOR_SENSOR, ATTR_DEVICE_TH_SENSOR, ATTR_DEVICE_MOTION_SENSOR],
     ),
@@ -89,18 +91,21 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up YoLink Sensor from a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][ATTR_COORDINATOR]
+    device_coordinators = hass.data[DOMAIN][config_entry.entry_id][ATTR_COORDINATORS]
-    sensor_devices = [
+    sensor_device_coordinators = [
-        device
+        device_coordinator
-        for device in coordinator.yl_devices
+        for device_coordinator in device_coordinators.values()
-        if device.device_type in SENSOR_DEVICE_TYPE
+        if device_coordinator.device.device_type in SENSOR_DEVICE_TYPE
     ]
     entities = []
-    for sensor_device in sensor_devices:
+    for sensor_device_coordinator in sensor_device_coordinators:
         for description in SENSOR_TYPES:
-            if description.exists_fn(sensor_device):
+            if description.exists_fn(sensor_device_coordinator.device):
                 entities.append(
-                    YoLinkSensorEntity(coordinator, description, sensor_device)
+                    YoLinkSensorEntity(
+                        sensor_device_coordinator,
+                        description,
+                    )
                 )
     async_add_entities(entities)
 
@@ -114,18 +119,21 @@ class YoLinkSensorEntity(YoLinkEntity, SensorEntity):
         self,
         coordinator: YoLinkCoordinator,
         description: YoLinkSensorEntityDescription,
-        device: YoLinkDevice,
     ) -> None:
         """Init YoLink Sensor."""
-        super().__init__(coordinator, device)
+        super().__init__(coordinator)
         self.entity_description = description
-        self._attr_unique_id = f"{device.device_id} {self.entity_description.key}"
+        self._attr_unique_id = (
-        self._attr_name = f"{device.device_name} ({self.entity_description.name})"
+            f"{coordinator.device.device_id} {self.entity_description.key}"
+        )
+        self._attr_name = (
+            f"{coordinator.device.device_name} ({self.entity_description.name})"
+        )
 
     @callback
     def update_entity_state(self, state: dict) -> None:
         """Update HA Entity State."""
         self._attr_native_value = self.entity_description.value(
-            state[self.entity_description.key]
+            state.get(self.entity_description.key)
         )
         self.async_write_ha_state()

homeassistant/components/yolink/siren.py

@@ -18,7 +18,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import ATTR_COORDINATOR, ATTR_DEVICE_SIREN, DOMAIN
+from .const import ATTR_COORDINATORS, ATTR_DEVICE_SIREN, DOMAIN
 from .coordinator import YoLinkCoordinator
 from .entity import YoLinkEntity
 
@@ -28,14 +28,14 @@ class YoLinkSirenEntityDescription(SirenEntityDescription):
     """YoLink SirenEntityDescription."""
 
     exists_fn: Callable[[YoLinkDevice], bool] = lambda _: True
-    value: Callable[[str], bool | None] = lambda _: None
+    value: Callable[[Any], bool | None] = lambda _: None
 
 
 DEVICE_TYPES: tuple[YoLinkSirenEntityDescription, ...] = (
     YoLinkSirenEntityDescription(
         key="state",
         name="State",
-        value=lambda value: value == "alert",
+        value=lambda value: value == "alert" if value is not None else None,
         exists_fn=lambda device: device.device_type in [ATTR_DEVICE_SIREN],
     ),
 )
@@ -49,16 +49,20 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up YoLink siren from a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][ATTR_COORDINATOR]
+    device_coordinators = hass.data[DOMAIN][config_entry.entry_id][ATTR_COORDINATORS]
-    devices = [
+    siren_device_coordinators = [
-        device for device in coordinator.yl_devices if device.device_type in DEVICE_TYPE
+        device_coordinator
+        for device_coordinator in device_coordinators.values()
+        if device_coordinator.device.device_type in DEVICE_TYPE
     ]
     entities = []
-    for device in devices:
+    for siren_device_coordinator in siren_device_coordinators:
         for description in DEVICE_TYPES:
-            if description.exists_fn(device):
+            if description.exists_fn(siren_device_coordinator.device):
                 entities.append(
-                    YoLinkSirenEntity(config_entry, coordinator, description, device)
+                    YoLinkSirenEntity(
+                        config_entry, siren_device_coordinator, description
+                    )
                 )
     async_add_entities(entities)
 
@@ -73,23 +77,26 @@ class YoLinkSirenEntity(YoLinkEntity, SirenEntity):
         config_entry: ConfigEntry,
         coordinator: YoLinkCoordinator,
         description: YoLinkSirenEntityDescription,
-        device: YoLinkDevice,
     ) -> None:
         """Init YoLink Siren."""
-        super().__init__(coordinator, device)
+        super().__init__(coordinator)
         self.config_entry = config_entry
         self.entity_description = description
-        self._attr_unique_id = f"{device.device_id} {self.entity_description.key}"
+        self._attr_unique_id = (
-        self._attr_name = f"{device.device_name} ({self.entity_description.name})"
+            f"{coordinator.device.device_id} {self.entity_description.key}"
+        )
+        self._attr_name = (
+            f"{coordinator.device.device_name} ({self.entity_description.name})"
+        )
         self._attr_supported_features = (
             SirenEntityFeature.TURN_ON | SirenEntityFeature.TURN_OFF
         )
 
     @callback
-    def update_entity_state(self, state: dict) -> None:
+    def update_entity_state(self, state: dict[str, Any]) -> None:
         """Update HA Entity State."""
         self._attr_is_on = self.entity_description.value(
-            state[self.entity_description.key]
+            state.get(self.entity_description.key)
         )
         self.async_write_ha_state()
 
@@ -97,7 +104,7 @@ class YoLinkSirenEntity(YoLinkEntity, SirenEntity):
         """Call setState api to change siren state."""
         try:
             # call_device_http_api will check result, fail by raise YoLinkClientError
-            await self.device.call_device_http_api(
+            await self.coordinator.device.call_device_http_api(
                 "setState", {"state": {"alarm": state}}
             )
         except YoLinkAuthFailError as yl_auth_err:

homeassistant/components/yolink/switch.py

@@ -18,7 +18,7 @@ from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from .const import ATTR_COORDINATOR, ATTR_DEVICE_OUTLET, DOMAIN
+from .const import ATTR_COORDINATORS, ATTR_DEVICE_OUTLET, DOMAIN
 from .coordinator import YoLinkCoordinator
 from .entity import YoLinkEntity
 
@@ -28,7 +28,7 @@ class YoLinkSwitchEntityDescription(SwitchEntityDescription):
     """YoLink SwitchEntityDescription."""
 
     exists_fn: Callable[[YoLinkDevice], bool] = lambda _: True
-    value: Callable[[str], bool | None] = lambda _: None
+    value: Callable[[Any], bool | None] = lambda _: None
 
 
 DEVICE_TYPES: tuple[YoLinkSwitchEntityDescription, ...] = (
@@ -36,7 +36,7 @@ DEVICE_TYPES: tuple[YoLinkSwitchEntityDescription, ...] = (
         key="state",
         device_class=SwitchDeviceClass.OUTLET,
         name="State",
-        value=lambda value: value == "open",
+        value=lambda value: value == "open" if value is not None else None,
         exists_fn=lambda device: device.device_type in [ATTR_DEVICE_OUTLET],
     ),
 )
@@ -50,16 +50,20 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up YoLink Sensor from a config entry."""
-    coordinator = hass.data[DOMAIN][config_entry.entry_id][ATTR_COORDINATOR]
+    device_coordinators = hass.data[DOMAIN][config_entry.entry_id][ATTR_COORDINATORS]
-    devices = [
+    switch_device_coordinators = [
-        device for device in coordinator.yl_devices if device.device_type in DEVICE_TYPE
+        device_coordinator
+        for device_coordinator in device_coordinators.values()
+        if device_coordinator.device.device_type in DEVICE_TYPE
     ]
     entities = []
-    for device in devices:
+    for switch_device_coordinator in switch_device_coordinators:
         for description in DEVICE_TYPES:
-            if description.exists_fn(device):
+            if description.exists_fn(switch_device_coordinator.device):
                 entities.append(
-                    YoLinkSwitchEntity(config_entry, coordinator, description, device)
+                    YoLinkSwitchEntity(
+                        config_entry, switch_device_coordinator, description
+                    )
                 )
     async_add_entities(entities)
 
@@ -74,20 +78,23 @@ class YoLinkSwitchEntity(YoLinkEntity, SwitchEntity):
         config_entry: ConfigEntry,
         coordinator: YoLinkCoordinator,
         description: YoLinkSwitchEntityDescription,
-        device: YoLinkDevice,
     ) -> None:
         """Init YoLink Outlet."""
-        super().__init__(coordinator, device)
+        super().__init__(coordinator)
         self.config_entry = config_entry
         self.entity_description = description
-        self._attr_unique_id = f"{device.device_id} {self.entity_description.key}"
+        self._attr_unique_id = (
-        self._attr_name = f"{device.device_name} ({self.entity_description.name})"
+            f"{coordinator.device.device_id} {self.entity_description.key}"
+        )
+        self._attr_name = (
+            f"{coordinator.device.device_name} ({self.entity_description.name})"
+        )
 
     @callback
-    def update_entity_state(self, state: dict) -> None:
+    def update_entity_state(self, state: dict[str, Any]) -> None:
         """Update HA Entity State."""
         self._attr_is_on = self.entity_description.value(
-            state[self.entity_description.key]
+            state.get(self.entity_description.key)
         )
         self.async_write_ha_state()
 
@@ -95,7 +102,9 @@ class YoLinkSwitchEntity(YoLinkEntity, SwitchEntity):
         """Call setState api to change outlet state."""
         try:
             # call_device_http_api will check result, fail by raise YoLinkClientError
-            await self.device.call_device_http_api("setState", {"state": state})
+            await self.coordinator.device.call_device_http_api(
+                "setState", {"state": state}
+            )
         except YoLinkAuthFailError as yl_auth_err:
             self.config_entry.async_start_reauth(self.hass)
             raise HomeAssistantError(yl_auth_err) from yl_auth_err