Coordinator refactor in Elgato (#87490)

homeassistant/components/elgato/__init__.py

@@ -1,59 +1,20 @@
 """Support for Elgato Lights."""
-from typing import NamedTuple
-
-from elgato import Elgato, ElgatoConnectionError, Info, State
-
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_HOST, CONF_PORT, Platform
+from homeassistant.const import Platform
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers.aiohttp_client import async_get_clientsession
-from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
 
-from .const import DOMAIN, LOGGER, SCAN_INTERVAL
+from .const import DOMAIN
+from .coordinator import ElgatoDataUpdateCoordinator
 
 PLATFORMS = [Platform.BUTTON, Platform.LIGHT]
 
 
-class HomeAssistantElgatoData(NamedTuple):
-    """Elgato data stored in the Home Assistant data object."""
-
-    coordinator: DataUpdateCoordinator[State]
-    client: Elgato
-    info: Info
-
-
 async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Elgato Light from a config entry."""
-    session = async_get_clientsession(hass)
-    elgato = Elgato(
-        entry.data[CONF_HOST],
-        port=entry.data[CONF_PORT],
-        session=session,
-    )
-
-    async def _async_update_data() -> State:
-        """Fetch Elgato data."""
-        try:
-            return await elgato.state()
-        except ElgatoConnectionError as err:
-            raise UpdateFailed(err) from err
-
-    coordinator: DataUpdateCoordinator[State] = DataUpdateCoordinator(
-        hass,
-        LOGGER,
-        name=f"{DOMAIN}_{entry.data[CONF_HOST]}",
-        update_interval=SCAN_INTERVAL,
-        update_method=_async_update_data,
-    )
+    coordinator = ElgatoDataUpdateCoordinator(hass, entry)
     await coordinator.async_config_entry_first_refresh()
 
-    info = await elgato.info()
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = HomeAssistantElgatoData(
-        client=elgato,
-        coordinator=coordinator,
-        info=info,
-    )
-
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = coordinator
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
     return True
@@ -62,8 +23,5 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
 async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Unload Elgato Light config entry."""
     if unload_ok := await hass.config_entries.async_unload_platforms(entry, PLATFORMS):
-        # Cleanup
         del hass.data[DOMAIN][entry.entry_id]
-        if not hass.data[DOMAIN]:
-            del hass.data[DOMAIN]
     return unload_ok

homeassistant/components/elgato/button.py

@@ -1,24 +1,19 @@
 """Support for Elgato button."""
 from __future__ import annotations
 
-import logging
-
-from elgato import Elgato, ElgatoError, Info
+from elgato import ElgatoError
 
 from homeassistant.components.button import ButtonEntity, ButtonEntityDescription
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_MAC
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity import EntityCategory
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
 
-from . import HomeAssistantElgatoData
 from .const import DOMAIN
+from .coordinator import ElgatoDataUpdateCoordinator
 from .entity import ElgatoEntity
 
-_LOGGER = logging.getLogger(__name__)
-
 
 async def async_setup_entry(
     hass: HomeAssistant,
@@ -26,30 +21,30 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Elgato button based on a config entry."""
-    data: HomeAssistantElgatoData = hass.data[DOMAIN][entry.entry_id]
-    async_add_entities(
-        [ElgatoIdentifyButton(data.client, data.info, entry.data.get(CONF_MAC))]
-    )
+    coordinator: ElgatoDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities([ElgatoIdentifyButton(coordinator)])
 
 
 class ElgatoIdentifyButton(ElgatoEntity, ButtonEntity):
     """Defines an Elgato identify button."""
 
-    def __init__(self, client: Elgato, info: Info, mac: str | None) -> None:
+    def __init__(self, coordinator: ElgatoDataUpdateCoordinator) -> None:
         """Initialize the button entity."""
-        super().__init__(client, info, mac)
+        super().__init__(coordinator=coordinator)
         self.entity_description = ButtonEntityDescription(
             key="identify",
             name="Identify",
             icon="mdi:help",
             entity_category=EntityCategory.CONFIG,
         )
-        self._attr_unique_id = f"{info.serial_number}_{self.entity_description.key}"
+        self._attr_unique_id = (
+            f"{coordinator.data.info.serial_number}_{self.entity_description.key}"
+        )
 
     async def async_press(self) -> None:
         """Identify the light, will make it blink."""
         try:
-            await self.client.identify()
+            await self.coordinator.client.identify()
         except ElgatoError as error:
             raise HomeAssistantError(
                 "An error occurred while identifying the Elgato Light"

homeassistant/components/elgato/coordinator.py

@@ -0,0 +1,53 @@
+"""DataUpdateCoordinator for Elgato."""
+from dataclasses import dataclass
+
+from elgato import Elgato, ElgatoConnectionError, Info, Settings, State
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_HOST, CONF_PORT
+from homeassistant.core import HomeAssistant
+from homeassistant.helpers.aiohttp_client import async_get_clientsession
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator, UpdateFailed
+
+from .const import DOMAIN, LOGGER, SCAN_INTERVAL
+
+
+@dataclass
+class ElgatoData:
+    """Elgato data stored in the DataUpdateCoordinator."""
+
+    info: Info
+    settings: Settings
+    state: State
+
+
+class ElgatoDataUpdateCoordinator(DataUpdateCoordinator[ElgatoData]):
+    """Class to manage fetching Elgato data."""
+
+    config_entry: ConfigEntry
+
+    def __init__(self, hass: HomeAssistant, entry: ConfigEntry) -> None:
+        """Initialize the coordinator."""
+        self.config_entry = entry
+        self.client = Elgato(
+            entry.data[CONF_HOST],
+            port=entry.data[CONF_PORT],
+            session=async_get_clientsession(hass),
+        )
+        super().__init__(
+            hass,
+            LOGGER,
+            name=f"{DOMAIN}_{entry.data[CONF_HOST]}",
+            update_interval=SCAN_INTERVAL,
+        )
+
+    async def _async_update_data(self) -> ElgatoData:
+        """Fetch data from the Elgato device."""
+        try:
+            return ElgatoData(
+                info=await self.client.info(),
+                settings=await self.client.settings(),
+                state=await self.client.state(),
+            )
+        except ElgatoConnectionError as err:
+            raise UpdateFailed(err) from err

homeassistant/components/elgato/diagnostics.py

@@ -6,16 +6,16 @@ from typing import Any
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.core import HomeAssistant
 
-from . import HomeAssistantElgatoData
 from .const import DOMAIN
+from .coordinator import ElgatoDataUpdateCoordinator
 
 
 async def async_get_config_entry_diagnostics(
     hass: HomeAssistant, entry: ConfigEntry
 ) -> dict[str, Any]:
     """Return diagnostics for a config entry."""
-    data: HomeAssistantElgatoData = hass.data[DOMAIN][entry.entry_id]
+    coordinator: ElgatoDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
     return {
-        "info": data.info.dict(),
-        "state": data.coordinator.data.dict(),
+        "info": coordinator.data.info.dict(),
+        "state": coordinator.data.state.dict(),
     }

homeassistant/components/elgato/entity.py

@@ -1,31 +1,32 @@
 """Base entity for the Elgato integration."""
 from __future__ import annotations
 
-from elgato import Elgato, Info
-
+from homeassistant.const import CONF_MAC
 from homeassistant.helpers.device_registry import CONNECTION_NETWORK_MAC, format_mac
-from homeassistant.helpers.entity import DeviceInfo, Entity
+from homeassistant.helpers.entity import DeviceInfo
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 
 from .const import DOMAIN
+from .coordinator import ElgatoDataUpdateCoordinator
 
 
-class ElgatoEntity(Entity):
+class ElgatoEntity(CoordinatorEntity[ElgatoDataUpdateCoordinator]):
     """Defines an Elgato entity."""
 
     _attr_has_entity_name = True
 
-    def __init__(self, client: Elgato, info: Info, mac: str | None) -> None:
+    def __init__(self, coordinator: ElgatoDataUpdateCoordinator) -> None:
         """Initialize an Elgato entity."""
-        self.client = client
+        super().__init__(coordinator=coordinator)
         self._attr_device_info = DeviceInfo(
-            identifiers={(DOMAIN, info.serial_number)},
+            identifiers={(DOMAIN, coordinator.data.info.serial_number)},
             manufacturer="Elgato",
-            model=info.product_name,
-            name=info.display_name,
-            sw_version=f"{info.firmware_version} ({info.firmware_build_number})",
-            hw_version=str(info.hardware_board_type),
+            model=coordinator.data.info.product_name,
+            name=coordinator.data.info.display_name,
+            sw_version=f"{coordinator.data.info.firmware_version} ({coordinator.data.info.firmware_build_number})",
+            hw_version=str(coordinator.data.info.hardware_board_type),
         )
-        if mac is not None:
+        if (mac := coordinator.config_entry.data.get(CONF_MAC)) is not None:
             self._attr_device_info["connections"] = {
                 (CONNECTION_NETWORK_MAC, format_mac(mac))
             }

homeassistant/components/elgato/light.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 
 from typing import Any
 
-from elgato import Elgato, ElgatoError, Info, Settings, State
+from elgato import ElgatoError
 
 from homeassistant.components.light import (
     ATTR_BRIGHTNESS,
@@ -13,20 +13,15 @@ from homeassistant.components.light import (
     LightEntity,
 )
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import CONF_MAC
 from homeassistant.core import HomeAssistant
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.entity_platform import (
     AddEntitiesCallback,
     async_get_current_platform,
 )
-from homeassistant.helpers.update_coordinator import (
-    CoordinatorEntity,
-    DataUpdateCoordinator,
-)
 
-from . import HomeAssistantElgatoData
 from .const import DOMAIN, SERVICE_IDENTIFY
+from .coordinator import ElgatoDataUpdateCoordinator
 from .entity import ElgatoEntity
 
 PARALLEL_UPDATES = 1
@@ -38,20 +33,8 @@ async def async_setup_entry(
     async_add_entities: AddEntitiesCallback,
 ) -> None:
     """Set up Elgato Light based on a config entry."""
-    data: HomeAssistantElgatoData = hass.data[DOMAIN][entry.entry_id]
-    settings = await data.client.settings()
-    async_add_entities(
-        [
-            ElgatoLight(
-                data.coordinator,
-                data.client,
-                data.info,
-                entry.data.get(CONF_MAC),
-                settings,
-            )
-        ],
-        True,
-    )
+    coordinator: ElgatoDataUpdateCoordinator = hass.data[DOMAIN][entry.entry_id]
+    async_add_entities([ElgatoLight(coordinator)])
 
     platform = async_get_current_platform()
     platform.async_register_entity_service(
@@ -61,30 +44,20 @@ async def async_setup_entry(
     )
 
 
-class ElgatoLight(
-    ElgatoEntity, CoordinatorEntity[DataUpdateCoordinator[State]], LightEntity
-):
+class ElgatoLight(ElgatoEntity, LightEntity):
     """Defines an Elgato Light."""
 
-    def __init__(
-        self,
-        coordinator: DataUpdateCoordinator[State],
-        client: Elgato,
-        info: Info,
-        mac: str | None,
-        settings: Settings,
-    ) -> None:
-        """Initialize Elgato Light."""
-        super().__init__(client, info, mac)
-        CoordinatorEntity.__init__(self, coordinator)
+    _attr_min_mireds = 143
+    _attr_max_mireds = 344
 
-        self._attr_min_mireds = 143
-        self._attr_max_mireds = 344
+    def __init__(self, coordinator: ElgatoDataUpdateCoordinator) -> None:
+        """Initialize Elgato Light."""
+        super().__init__(coordinator)
         self._attr_supported_color_modes = {ColorMode.COLOR_TEMP}
-        self._attr_unique_id = info.serial_number
+        self._attr_unique_id = coordinator.data.info.serial_number
 
         # Elgato Light supporting color, have a different temperature range
-        if settings.power_on_hue is not None:
+        if self.coordinator.data.settings.power_on_hue is not None:
             self._attr_supported_color_modes = {ColorMode.COLOR_TEMP, ColorMode.HS}
             self._attr_min_mireds = 153
             self._attr_max_mireds = 285
@@ -92,17 +65,17 @@ class ElgatoLight(
     @property
     def brightness(self) -> int | None:
         """Return the brightness of this light between 1..255."""
-        return round((self.coordinator.data.brightness * 255) / 100)
+        return round((self.coordinator.data.state.brightness * 255) / 100)
 
     @property
     def color_temp(self) -> int | None:
         """Return the CT color value in mireds."""
-        return self.coordinator.data.temperature
+        return self.coordinator.data.state.temperature
 
     @property
     def color_mode(self) -> str | None:
         """Return the color mode of the light."""
-        if self.coordinator.data.hue is not None:
+        if self.coordinator.data.state.hue is not None:
             return ColorMode.HS
 
         return ColorMode.COLOR_TEMP
@@ -110,17 +83,20 @@ class ElgatoLight(
     @property
     def hs_color(self) -> tuple[float, float] | None:
         """Return the hue and saturation color value [float, float]."""
-        return (self.coordinator.data.hue or 0, self.coordinator.data.saturation or 0)
+        return (
+            self.coordinator.data.state.hue or 0,
+            self.coordinator.data.state.saturation or 0,
+        )
 
     @property
     def is_on(self) -> bool:
         """Return the state of the light."""
-        return self.coordinator.data.on
+        return self.coordinator.data.state.on
 
     async def async_turn_off(self, **kwargs: Any) -> None:
         """Turn off the light."""
         try:
-            await self.client.light(on=False)
+            await self.coordinator.client.light(on=False)
         except ElgatoError as error:
             raise HomeAssistantError(
                 "An error occurred while updating the Elgato Light"
@@ -155,7 +131,7 @@ class ElgatoLight(
             temperature = self.color_temp
 
         try:
-            await self.client.light(
+            await self.coordinator.client.light(
                 on=True,
                 brightness=brightness,
                 hue=hue,
@@ -172,7 +148,7 @@ class ElgatoLight(
     async def async_identify(self) -> None:
         """Identify the light, will make it blink."""
         try:
-            await self.client.identify()
+            await self.coordinator.client.identify()
         except ElgatoError as error:
             raise HomeAssistantError(
                 "An error occurred while identifying the Elgato Light"

tests/components/elgato/conftest.py

@@ -65,7 +65,9 @@ def mock_elgato(request: pytest.FixtureRequest) -> Generator[None, MagicMock, No
     if hasattr(request, "param") and request.param:
         variant = request.param
 
-    with patch("homeassistant.components.elgato.Elgato", autospec=True) as elgato_mock:
+    with patch(
+        "homeassistant.components.elgato.coordinator.Elgato", autospec=True
+    ) as elgato_mock:
         elgato = elgato_mock.return_value
         elgato.info.return_value = Info.parse_raw(load_fixture("info.json", DOMAIN))
         elgato.state.return_value = State.parse_raw(