Incorporate GroupManager into HEOS Coordinator (home-assistant#136462)

* Incorporate GroupManager

* Update quality scale

* Fix group params

* Revert quality scale change

* Rename varaible

* Move group action implementaton out of coordinator

* Fix get_group_members hass access

* entity -> entity_id

homeassistant/components/heos/__init__.py

@@ -2,27 +2,17 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
 from datetime import timedelta
-import logging
-from typing import Any
-
-from pyheos import Heos, HeosError, HeosPlayer, const as heos_const
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import Platform
-from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
-from homeassistant.exceptions import HomeAssistantError
+from homeassistant.core import HomeAssistant
 from homeassistant.helpers import device_registry as dr
 import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.dispatcher import (
-    async_dispatcher_connect,
-    async_dispatcher_send,
-)
 from homeassistant.helpers.typing import ConfigType
 
 from . import services
-from .const import DOMAIN, SIGNAL_HEOS_PLAYER_ADDED, SIGNAL_HEOS_UPDATED
+from .const import DOMAIN
 from .coordinator import HeosCoordinator
 
 PLATFORMS = [Platform.MEDIA_PLAYER]
@@ -31,19 +21,7 @@
 
 CONFIG_SCHEMA = cv.empty_config_schema(DOMAIN)
 
-_LOGGER = logging.getLogger(__name__)
-
-
-@dataclass
-class HeosRuntimeData:
-    """Runtime data and coordinators for HEOS config entries."""
-
-    coordinator: HeosCoordinator
-    group_manager: GroupManager
-    players: dict[int, HeosPlayer]
-
-
-type HeosConfigEntry = ConfigEntry[HeosRuntimeData]
+type HeosConfigEntry = ConfigEntry[HeosCoordinator]
 
 
 async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
@@ -72,16 +50,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: HeosConfigEntry) -> bool
 
     coordinator = HeosCoordinator(hass, entry)
     await coordinator.async_setup()
-    # Preserve existing logic until migrated into coordinator
-    controller = coordinator.heos
-    players = controller.players
-
-    group_manager = GroupManager(hass, controller, players)
-
-    entry.runtime_data = HeosRuntimeData(coordinator, group_manager, players)
-
-    group_manager.connect_update()
-    entry.async_on_unload(group_manager.disconnect_update)
+    entry.runtime_data = coordinator
 
     await hass.config_entries.async_forward_entry_setups(entry, PLATFORMS)
 
@@ -91,130 +60,3 @@ async def async_setup_entry(hass: HomeAssistant, entry: HeosConfigEntry) -> bool
 async def async_unload_entry(hass: HomeAssistant, entry: HeosConfigEntry) -> bool:
     """Unload a config entry."""
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
-
-
-class GroupManager:
-    """Class that manages HEOS groups."""
-
-    def __init__(
-        self, hass: HomeAssistant, controller: Heos, players: dict[int, HeosPlayer]
-    ) -> None:
-        """Init group manager."""
-        self._hass = hass
-        self._group_membership: dict[str, list[str]] = {}
-        self._disconnect_player_added = None
-        self._initialized = False
-        self.controller = controller
-        self.players = players
-        self.entity_id_map: dict[int, str] = {}
-
-    def _get_entity_id_to_player_id_map(self) -> dict:
-        """Return mapping of all HeosMediaPlayer entity_ids to player_ids."""
-        return {v: k for k, v in self.entity_id_map.items()}
-
-    async def async_get_group_membership(self) -> dict[str, list[str]]:
-        """Return all group members for each player as entity_ids."""
-        group_info_by_entity_id: dict[str, list[str]] = {
-            player_entity_id: []
-            for player_entity_id in self._get_entity_id_to_player_id_map()
-        }
-
-        try:
-            groups = await self.controller.get_groups()
-        except HeosError as err:
-            _LOGGER.error("Unable to get HEOS group info: %s", err)
-            return group_info_by_entity_id
-
-        player_id_to_entity_id_map = self.entity_id_map
-        for group in groups.values():
-            leader_entity_id = player_id_to_entity_id_map.get(group.lead_player_id)
-            member_entity_ids = [
-                player_id_to_entity_id_map[member]
-                for member in group.member_player_ids
-                if member in player_id_to_entity_id_map
-            ]
-            # Make sure the group leader is always the first element
-            group_info = [leader_entity_id, *member_entity_ids]
-            if leader_entity_id:
-                group_info_by_entity_id[leader_entity_id] = group_info  # type: ignore[assignment]
-                for member_entity_id in member_entity_ids:
-                    group_info_by_entity_id[member_entity_id] = group_info  # type: ignore[assignment]
-
-        return group_info_by_entity_id
-
-    async def async_join_players(
-        self, leader_id: int, member_entity_ids: list[str]
-    ) -> None:
-        """Create a group a group leader and member players."""
-        # Resolve HEOS player_id for each member entity_id
-        entity_id_to_player_id_map = self._get_entity_id_to_player_id_map()
-        member_ids: list[int] = []
-        for member in member_entity_ids:
-            member_id = entity_id_to_player_id_map.get(member)
-            if not member_id:
-                raise HomeAssistantError(
-                    f"The group member {member} could not be resolved to a HEOS player."
-                )
-            member_ids.append(member_id)
-
-        await self.controller.create_group(leader_id, member_ids)
-
-    async def async_unjoin_player(self, player_id: int):
-        """Remove `player_entity_id` from any group."""
-        await self.controller.create_group(player_id, [])
-
-    async def async_update_groups(self) -> None:
-        """Update the group membership from the controller."""
-        if groups := await self.async_get_group_membership():
-            self._group_membership = groups
-            _LOGGER.debug("Groups updated due to change event")
-            # Let players know to update
-            async_dispatcher_send(self._hass, SIGNAL_HEOS_UPDATED)
-        else:
-            _LOGGER.debug("Groups empty")
-
-    @callback
-    def connect_update(self):
-        """Connect listener for when groups change and signal player update."""
-
-        async def _on_controller_event(event: str, data: Any | None) -> None:
-            if event == heos_const.EVENT_GROUPS_CHANGED:
-                await self.async_update_groups()
-
-        self.controller.add_on_controller_event(_on_controller_event)
-        self.controller.add_on_connected(self.async_update_groups)
-
-        # When adding a new HEOS player we need to update the groups.
-        async def _async_handle_player_added():
-            # Avoid calling async_update_groups when the entity_id map has not been
-            # fully populated yet. This may only happen during early startup.
-            if len(self.players) <= len(self.entity_id_map) and not self._initialized:
-                self._initialized = True
-                await self.async_update_groups()
-
-        self._disconnect_player_added = async_dispatcher_connect(
-            self._hass, SIGNAL_HEOS_PLAYER_ADDED, _async_handle_player_added
-        )
-
-    @callback
-    def disconnect_update(self):
-        """Disconnect the listeners."""
-        if self._disconnect_player_added:
-            self._disconnect_player_added()
-            self._disconnect_player_added = None
-
-    @callback
-    def register_media_player(self, player_id: int, entity_id: str) -> CALLBACK_TYPE:
-        """Register a media player player_id with it's entity_id so it can be resolved later."""
-        self.entity_id_map[player_id] = entity_id
-        return lambda: self.unregister_media_player(player_id)
-
-    @callback
-    def unregister_media_player(self, player_id) -> None:
-        """Remove a media player player_id from the entity_id map."""
-        self.entity_id_map.pop(player_id, None)
-
-    @property
-    def group_membership(self):
-        """Provide access to group members for player entities."""
-        return self._group_membership

homeassistant/components/heos/config_flow.py

@@ -188,9 +188,7 @@ async def async_step_reauth_confirm(
         entry: HeosConfigEntry = self._get_reauth_entry()
         if user_input is not None:
             assert entry.state is ConfigEntryState.LOADED
-            if await _validate_auth(
-                user_input, entry.runtime_data.coordinator.heos, errors
-            ):
+            if await _validate_auth(user_input, entry.runtime_data.heos, errors):
                 return self.async_update_reload_and_abort(entry, options=user_input)
 
         return self.async_show_form(
@@ -212,9 +210,7 @@ async def async_step_init(
         errors: dict[str, str] = {}
         if user_input is not None:
             entry: HeosConfigEntry = self.config_entry
-            if await _validate_auth(
-                user_input, entry.runtime_data.coordinator.heos, errors
-            ):
+            if await _validate_auth(user_input, entry.runtime_data.heos, errors):
                 return self.async_create_entry(data=user_input)
 
         return self.async_show_form(

homeassistant/components/heos/const.py

@@ -5,5 +5,3 @@
 DOMAIN = "heos"
 SERVICE_SIGN_IN = "sign_in"
 SERVICE_SIGN_OUT = "sign_out"
-SIGNAL_HEOS_PLAYER_ADDED = "heos_player_added"
-SIGNAL_HEOS_UPDATED = "heos_updated"

homeassistant/components/heos/coordinator.py

@@ -5,6 +5,7 @@
 entities to update. Entities subscribe to entity-specific updates within the entity class itself.
 """
 
+from collections.abc import Callable
 from datetime import datetime, timedelta
 import logging
 
@@ -81,6 +82,7 @@ async def async_setup(self) -> None:
                 "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_groups()
         await self._async_update_sources()
         # Attach event callbacks
         self.heos.add_on_disconnected(self._async_on_disconnected)
@@ -93,6 +95,13 @@ async def async_shutdown(self) -> None:
         await self.heos.disconnect()
         await super().async_shutdown()
 
+    def async_add_listener(self, update_callback, context=None) -> Callable[[], None]:
+        """Add a listener for the coordinator."""
+        remove_listener = super().async_add_listener(update_callback, context)
+        # Update entities so group_member entity_ids fully populate.
+        self.async_update_listeners()
+        return remove_listener
+
     async def _async_on_auth_failure(self) -> None:
         """Handle when the user credentials are no longer valid."""
         assert self.config_entry is not None
@@ -118,6 +127,8 @@ async def _async_on_controller_event(
             assert data is not None
             if data.updated_player_ids:
                 self._async_update_player_ids(data.updated_player_ids)
+        elif event == const.EVENT_GROUPS_CHANGED:
+            await self._async_update_players()
         elif (
             event in (const.EVENT_SOURCES_CHANGED, const.EVENT_USER_CHANGED)
             and not self._update_sources_pending
@@ -176,6 +187,13 @@ def _async_update_player_ids(self, updated_player_ids: dict[int, int]) -> None:
                 )
                 _LOGGER.debug("Updated entity %s unique id to %s", entity_id, new_id)
 
+    async def _async_update_groups(self) -> None:
+        """Update group information."""
+        try:
+            await self.heos.get_groups(refresh=True)
+        except HeosError as error:
+            _LOGGER.error("Unable to retrieve groups: %s", error)
+
     async def _async_update_sources(self) -> None:
         """Build source list for entities."""
         self._source_list.clear()