Implement Coordinator for HEOS (initial plumbing) (#136205)

homeassistant/components/heos/__init__.py

@@ -9,10 +9,8 @@ import logging
 from typing import Any
 
 from pyheos import (
-    Credentials,
     Heos,
     HeosError,
-    HeosOptions,
     HeosPlayer,
     PlayerUpdateResult,
     SignalHeosEvent,
@@ -20,19 +18,9 @@ from pyheos import (
 )
 
 from homeassistant.config_entries import ConfigEntry
-from homeassistant.const import (
+from homeassistant.const import Platform
-    CONF_HOST,
-    CONF_PASSWORD,
-    CONF_USERNAME,
-    EVENT_HOMEASSISTANT_STOP,
-    Platform,
-)
 from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
-from homeassistant.exceptions import (
+from homeassistant.exceptions import HomeAssistantError, ServiceValidationError
-    ConfigEntryNotReady,
-    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 @@ from .const import (
     SIGNAL_HEOS_PLAYER_ADDED,
     SIGNAL_HEOS_UPDATED,
 )
+from .coordinator import HeosCoordinator
 
 PLATFORMS = [Platform.MEDIA_PLAYER]
 
@@ -64,6 +53,7 @@ _LOGGER = logging.getLogger(__name__)
 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]
+    coordinator = HeosCoordinator(hass, entry)
-    credentials: Credentials | None = None
+    await coordinator.async_setup()
-    if entry.options:
+    # Preserve existing logic until migrated into coordinator
-        credentials = Credentials(
+    controller = coordinator.heos
-            entry.options[CONF_USERNAME], entry.options[CONF_PASSWORD]
+    players = controller.players
-        )
+    favorites = coordinator.favorites
-
+    inputs = coordinator.inputs
-    # 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
 
     controller_manager = ControllerManager(hass, controller)
     await controller_manager.connect_listeners()
@@ -164,7 +104,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: HeosConfigEntry) -> bool
     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 async_setup_entry(hass: HomeAssistant, entry: HeosConfigEntry) -> bool
 
 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,7 +29,7 @@ from homeassistant.components.media_player import (
     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 (
@@ -37,10 +37,12 @@ from homeassistant.helpers.dispatcher import (
     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,11 +131,10 @@ def catch_action_error[**_P](
     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
@@ -138,6 +142,7 @@ class HeosMediaPlayer(MediaPlayerEntity):
 
     def __init__(
         self,
+        coordinator: HeosCoordinator,
         player: HeosPlayer,
         source_manager: SourceManager,
         group_manager: GroupManager,
@@ -159,16 +164,34 @@ class HeosMediaPlayer(MediaPlayerEntity):
             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:
+    @callback
-        """Handle sources changed."""
+    def _handle_coordinator_update(self) -> None:
-        await self.async_update_ha_state(True)
+        """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."""
@@ -176,7 +199,9 @@ class HeosMediaPlayer(MediaPlayerEntity):
         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 @@ class HeosMediaPlayer(MediaPlayerEntity):
             )
         )
         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 @@ class HeosMediaPlayer(MediaPlayerEntity):
         """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

tests/components/heos/test_init.py

@@ -30,7 +30,7 @@ async def test_async_setup_entry_loads_platforms(
     """Test load connects to heos, retrieves players, and loads platforms."""
     config_entry.add_to_hass(hass)
     assert await hass.config_entries.async_setup(config_entry.entry_id)
-    assert config_entry.state == ConfigEntryState.LOADED
+    assert config_entry.state is ConfigEntryState.LOADED
     assert hass.states.get("media_player.test_player") is not None
     assert controller.connect.call_count == 1
     assert controller.get_players.call_count == 1
@@ -116,24 +116,41 @@ async def test_async_setup_entry_connect_failure(
     config_entry.add_to_hass(hass)
     controller.connect.side_effect = HeosError()
     assert not await hass.config_entries.async_setup(config_entry.entry_id)
-    assert config_entry.state == ConfigEntryState.SETUP_RETRY
     assert controller.connect.call_count == 1
     assert controller.disconnect.call_count == 1
-    controller.connect.reset_mock()
+    assert config_entry.state is ConfigEntryState.SETUP_RETRY
-    controller.disconnect.reset_mock()
 
 
 async def test_async_setup_entry_player_failure(
     hass: HomeAssistant, config_entry: MockConfigEntry, controller: Heos
 ) -> None:
-    """Failure to retrieve players/sources raises ConfigEntryNotReady."""
+    """Failure to retrieve players raises ConfigEntryNotReady."""
     config_entry.add_to_hass(hass)
     controller.get_players.side_effect = HeosError()
     assert not await hass.config_entries.async_setup(config_entry.entry_id)
     assert controller.connect.call_count == 1
     assert controller.disconnect.call_count == 1
-    controller.connect.reset_mock()
+    assert config_entry.state is ConfigEntryState.SETUP_RETRY
-    controller.disconnect.reset_mock()
+
+
+async def test_async_setup_entry_favorites_failure(
+    hass: HomeAssistant, config_entry: MockConfigEntry, controller: Heos
+) -> None:
+    """Failure to retrieve favorites loads."""
+    config_entry.add_to_hass(hass)
+    controller.get_favorites.side_effect = HeosError()
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    assert config_entry.state is ConfigEntryState.LOADED
+
+
+async def test_async_setup_entry_inputs_failure(
+    hass: HomeAssistant, config_entry: MockConfigEntry, controller: Heos
+) -> None:
+    """Failure to retrieve inputs loads."""
+    config_entry.add_to_hass(hass)
+    controller.get_input_sources.side_effect = HeosError()
+    assert await hass.config_entries.async_setup(config_entry.entry_id)
+    assert config_entry.state is ConfigEntryState.LOADED
 
 
 async def test_unload_entry(