Implement Coordinator for HEOS (initial plumbing)

homeassistant/components/heos/__init__.py

@@ -9,30 +9,18 @@
 from typing import Any
 
 from pyheos import (
-    Credentials,
     Heos,
     HeosError,
-    HeosOptions,
     HeosPlayer,
     PlayerUpdateResult,
     SignalHeosEvent,
     const as heos_const,
 )
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
-    CONF_HOST,
-    CONF_PASSWORD,
-    CONF_USERNAME,
-    EVENT_HOMEASSISTANT_STOP,
-    Platform,
-)
+from homeassistant.const import Platform
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
-from homeassistant.exceptions import (
-    ConfigEntryNotReady,
-    HomeAssistantError,
-    ServiceValidationError,
-)
+from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.dispatcher import (
@@ -50,6 +38,7 @@
     SIGNAL_HEOS_PLAYER_ADDED,
     SIGNAL_HEOS_UPDATED,
 )
+from .coordinator import HeosCoordinator
 
 PLATFORMS = [Platform.MEDIA_PLAYER]
 
@@ -64,6 +53,7 @@
 class HeosRuntimeData:
     """Runtime data and coordinators for HEOS config entries."""
 
+    coordinator: HeosCoordinator
     controller_manager: ControllerManager
     group_manager: GroupManager
     source_manager: SourceManager
@@ -97,63 +87,13 @@ async def async_setup_entry(hass: HomeAssistant, entry: HeosConfigEntry) -> bool
                 )
             break
 
-    host = entry.data[CONF_HOST]
-    credentials: Credentials | None = None
-    if entry.options:
-        credentials = Credentials(
-            entry.options[CONF_USERNAME], entry.options[CONF_PASSWORD]
-        )
-
-    # Setting all_progress_events=False ensures that we only receive a
-    # media position update upon start of playback or when media changes
-    controller = Heos(
-        HeosOptions(
-            host,
-            all_progress_events=False,
-            auto_reconnect=True,
-            credentials=credentials,
-        )
-    )
-
-    # Auth failure handler must be added before connecting to the host, otherwise
-    # the event will be missed when login fails during connection.
-    async def auth_failure() -> None:
-        """Handle authentication failure."""
-        entry.async_start_reauth(hass)
-
-    entry.async_on_unload(controller.add_on_user_credentials_invalid(auth_failure))
-
-    try:
-        # Auto reconnect only operates if initial connection was successful.
-        await controller.connect()
-    except HeosError as error:
-        await controller.disconnect()
-        _LOGGER.debug("Unable to connect to controller %s: %s", host, error)
-        raise ConfigEntryNotReady from error
-
-    # Disconnect when shutting down
-    async def disconnect_controller(event):
-        await controller.disconnect()
-
-    entry.async_on_unload(
-        hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, disconnect_controller)
-    )
-
-    # Get players and sources
-    try:
-        players = await controller.get_players()
-        favorites = {}
-        if controller.is_signed_in:
-            favorites = await controller.get_favorites()
-        else:
-            _LOGGER.warning(
-                "The HEOS System is not logged in: Enter credentials in the integration options to access favorites and streaming services"
-            )
-        inputs = await controller.get_input_sources()
-    except HeosError as error:
-        await controller.disconnect()
-        _LOGGER.debug("Unable to retrieve players and sources: %s", error)
-        raise ConfigEntryNotReady from error
+    coordinator = HeosCoordinator(hass, entry)
+    await coordinator.async_setup()
+    # Preserve existing logic until migrated into coordinator
+    controller = coordinator.heos
+    players = controller.players
+    favorites = coordinator.favorites
+    inputs = coordinator.inputs
 
     controller_manager = ControllerManager(hass, controller)
     await controller_manager.connect_listeners()
@@ -164,7 +104,7 @@ async def disconnect_controller(event):
     group_manager = GroupManager(hass, controller, players)
 
     entry.runtime_data = HeosRuntimeData(
-        controller_manager, group_manager, source_manager, players
+        coordinator, controller_manager, group_manager, source_manager, players
     )
 
     group_manager.connect_update()
@@ -177,7 +117,6 @@ async def disconnect_controller(event):
 
 async def async_unload_entry(hass: HomeAssistant, entry: HeosConfigEntry) -> bool:
     """Unload a config entry."""
-    await entry.runtime_data.controller_manager.disconnect()
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
 
 

homeassistant/components/heos/coordinator.py

@@ -0,0 +1,93 @@
+"""HEOS integration coordinator.
+
+Control of all HEOS devices is through connection to a single device. Data is pushed through events.
+The coordinator is responsible for refreshing data in response to system-wide events and notifying
+entities to update. Entities subscribe to entity-specific updates within the entity class itself.
+"""
+
+import logging
+
+from pyheos import Credentials, Heos, HeosError, HeosOptions, MediaItem
+
+from homeassistant.config_entries import ConfigEntry
+from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME
+from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import ConfigEntryNotReady
+from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+
+from . import DOMAIN
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class HeosCoordinator(DataUpdateCoordinator[None]):
+    """Define the HEOS integration coordinator."""
+
+    def __init__(self, hass: HomeAssistant, config_entry: ConfigEntry) -> None:
+        """Set up the coordinator and set in config_entry."""
+        self.host: str = config_entry.data[CONF_HOST]
+        credentials: Credentials | None = None
+        if config_entry.options:
+            credentials = Credentials(
+                config_entry.options[CONF_USERNAME], config_entry.options[CONF_PASSWORD]
+            )
+        # Setting all_progress_events=False ensures that we only receive a
+        # media position update upon start of playback or when media changes
+        self.heos = Heos(
+            HeosOptions(
+                self.host,
+                all_progress_events=False,
+                auto_reconnect=True,
+                credentials=credentials,
+            )
+        )
+        self.favorites: dict[int, MediaItem] = {}
+        self.inputs: list[MediaItem] = []
+        super().__init__(hass, _LOGGER, config_entry=config_entry, name=DOMAIN)
+
+    async def async_setup(self) -> None:
+        """Set up the coordinator; connect to the host; and retrieve initial data."""
+        # Add before connect as it may occur during initial connection
+        self.heos.add_on_user_credentials_invalid(self._async_on_auth_failure)
+        # Connect to the device
+        try:
+            await self.heos.connect()
+        except HeosError as error:
+            raise ConfigEntryNotReady from error
+        # Load players
+        try:
+            await self.heos.get_players()
+        except HeosError as error:
+            raise ConfigEntryNotReady from error
+
+        if not self.heos.is_signed_in:
+            _LOGGER.warning(
+                "The HEOS System is not logged in: Enter credentials in the integration options to access favorites and streaming services"
+            )
+        # Retrieve initial data
+        await self._async_update_sources()
+
+    async def async_shutdown(self) -> None:
+        """Disconnect all callbacks and disconnect from the device."""
+        self.heos.dispatcher.disconnect_all()  # Removes all connected through heos.add_on_* and player.add_on_*
+        await self.heos.disconnect()
+        await super().async_shutdown()
+
+    async def _async_on_auth_failure(self) -> None:
+        """Handle when the user credentials are no longer valid."""
+        assert self.config_entry is not None
+        self.config_entry.async_start_reauth(self.hass)
+
+    async def _async_update_sources(self) -> None:
+        """Build source list for entities."""
+        # Get favorites only if reportedly signed in.
+        if self.heos.is_signed_in:
+            try:
+                self.favorites = await self.heos.get_favorites()
+            except HeosError as error:
+                _LOGGER.error("Unable to retrieve favorites: %s", error)
+        # Get input sources (across all devices in the HEOS system)
+        try:
+            self.inputs = await self.heos.get_input_sources()
+        except HeosError as error:
+            _LOGGER.error("Unable to retrieve input sources: %s", error)

homeassistant/components/heos/media_player.py

@@ -29,18 +29,20 @@
     RepeatMode,
     async_process_play_media_url,
 )
-from homeassistant.core import HomeAssistant
+from homeassistant.core import HomeAssistant, callback
 from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers.device_registry import DeviceInfo
 from homeassistant.helpers.dispatcher import (
     async_dispatcher_connect,
     async_dispatcher_send,
 )
 from homeassistant.helpers.entity_platform import AddEntitiesCallback
+from homeassistant.helpers.update_coordinator import CoordinatorEntity
 from homeassistant.util.dt import utcnow
 
 from . import GroupManager, HeosConfigEntry, SourceManager
 from .const import DOMAIN as HEOS_DOMAIN, SIGNAL_HEOS_PLAYER_ADDED, SIGNAL_HEOS_UPDATED
+from .coordinator import HeosCoordinator
 
 PARALLEL_UPDATES = 0
 
@@ -93,11 +95,14 @@ async def async_setup_entry(
     players = entry.runtime_data.players
     devices = [
         HeosMediaPlayer(
-            player, entry.runtime_data.source_manager, entry.runtime_data.group_manager
+            entry.runtime_data.coordinator,
+            player,
+            entry.runtime_data.source_manager,
+            entry.runtime_data.group_manager,
         )
         for player in players.values()
     ]
-    async_add_entities(devices, True)
+    async_add_entities(devices)
 
 
 type _FuncType[**_P] = Callable[_P, Awaitable[Any]]
@@ -126,18 +131,18 @@ async def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
     return decorator
 
 
-class HeosMediaPlayer(MediaPlayerEntity):
+class HeosMediaPlayer(CoordinatorEntity[HeosCoordinator], MediaPlayerEntity):
     """The HEOS player."""
 
     _attr_media_content_type = MediaType.MUSIC
-    _attr_should_poll = False
     _attr_supported_features = BASE_SUPPORTED_FEATURES
     _attr_media_image_remotely_accessible = True
     _attr_has_entity_name = True
     _attr_name = None
 
     def __init__(
         self,
+        coordinator: HeosCoordinator,
         player: HeosPlayer,
         source_manager: SourceManager,
         group_manager: GroupManager,
@@ -159,24 +164,44 @@ def __init__(
             serial_number=player.serial,  # Only available for some models
             sw_version=player.version,
         )
+        self._update_attributes()
+        super().__init__(coordinator, context=player.player_id)
 
     async def _player_update(self, event):
         """Handle player attribute updated."""
         if event == heos_const.EVENT_PLAYER_NOW_PLAYING_PROGRESS:
             self._media_position_updated_at = utcnow()
-        await self.async_update_ha_state(True)
+        self._handle_coordinator_update()
 
-    async def _heos_updated(self) -> None:
-        """Handle sources changed."""
-        await self.async_update_ha_state(True)
+    @callback
+    def _handle_coordinator_update(self) -> None:
+        """Handle updated data from the coordinator."""
+        self._update_attributes()
+        super()._handle_coordinator_update()
+
+    def _update_attributes(self) -> None:
+        """Update core attributes of the media player."""
+        self._attr_repeat = HEOS_HA_REPEAT_TYPE_MAP[self._player.repeat]
+        controls = self._player.now_playing_media.supported_controls
+        current_support = [CONTROL_TO_SUPPORT[control] for control in controls]
+        self._attr_supported_features = reduce(
+            ior, current_support, BASE_SUPPORTED_FEATURES
+        )
+        if self.support_next_track and self.support_previous_track:
+            self._attr_supported_features |= (
+                MediaPlayerEntityFeature.REPEAT_SET
+                | MediaPlayerEntityFeature.SHUFFLE_SET
+            )
 
     async def async_added_to_hass(self) -> None:
         """Device added to hass."""
         # Update state when attributes of the player change
         self.async_on_remove(self._player.add_on_player_event(self._player_update))
         # Update state when heos changes
         self.async_on_remove(
-            async_dispatcher_connect(self.hass, SIGNAL_HEOS_UPDATED, self._heos_updated)
+            async_dispatcher_connect(
+                self.hass, SIGNAL_HEOS_UPDATED, self._handle_coordinator_update
+            )
         )
         # Register this player's entity_id so it can be resolved by the group manager
         self.async_on_remove(
@@ -185,6 +210,7 @@ async def async_added_to_hass(self) -> None:
             )
         )
         async_dispatcher_send(self.hass, SIGNAL_HEOS_PLAYER_ADDED)
+        await super().async_added_to_hass()
 
     @catch_action_error("clear playlist")
     async def async_clear_playlist(self) -> None:
@@ -315,20 +341,6 @@ async def async_set_volume_level(self, volume: float) -> None:
         """Set volume level, range 0..1."""
         await self._player.set_volume(int(volume * 100))
 
-    async def async_update(self) -> None:
-        """Update supported features of the player."""
-        self._attr_repeat = HEOS_HA_REPEAT_TYPE_MAP[self._player.repeat]
-        controls = self._player.now_playing_media.supported_controls
-        current_support = [CONTROL_TO_SUPPORT[control] for control in controls]
-        self._attr_supported_features = reduce(
-            ior, current_support, BASE_SUPPORTED_FEATURES
-        )
-        if self.support_next_track and self.support_previous_track:
-            self._attr_supported_features |= (
-                MediaPlayerEntityFeature.REPEAT_SET
-                | MediaPlayerEntityFeature.SHUFFLE_SET
-            )
-
     @catch_action_error("unjoin player")
     async def async_unjoin_player(self) -> None:
         """Remove this player from any group."""

tests/components/heos/conftest.py

@@ -91,7 +91,7 @@ async def controller_fixture(
     new_mock = Mock(return_value=mock_heos)
     mock_heos.new_mock = new_mock
     with (
-        patch("homeassistant.components.heos.Heos", new=new_mock),
+        patch("homeassistant.components.heos.coordinator.Heos", new=new_mock),
         patch("homeassistant.components.heos.config_flow.Heos", new=new_mock),
     ):
         yield mock_heos