Improve HEOS group handling (#132213)

* Move register method to GroupManager

* Remove GroupManager mapping when entity removed

* Add test for when unloaded

* Error when group member not found

* Use entity registery to remove entity

* Update tests per feedback

homeassistant/components/heos/__init__.py

@@ -11,7 +11,7 @@ from pyheos import Heos, HeosError, HeosPlayer, const as heos_const
 
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import CONF_HOST, EVENT_HOMEASSISTANT_STOP, Platform
-from homeassistant.core import HomeAssistant, callback
+from homeassistant.core import CALLBACK_TYPE, HomeAssistant, callback
 from homeassistant.exceptions import ConfigEntryNotReady, HomeAssistantError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.dispatcher import (
@@ -259,21 +259,19 @@ class GroupManager:
         return group_info_by_entity_id
 
     async def async_join_players(
-        self, leader_entity_id: str, member_entity_ids: list[str]
+        self, leader_id: int, leader_entity_id: str, 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()
-        leader_id = entity_id_to_player_id_map.get(leader_entity_id)
-        if not leader_id:
+        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 leader {leader_entity_id} could not be resolved to a HEOS"
-                " player."
+                    f"The group member {member} could not be resolved to a HEOS player."
                 )
-        member_ids = [
-            entity_id_to_player_id_map[member]
-            for member in member_entity_ids
-            if member in entity_id_to_player_id_map
-        ]
+            member_ids.append(member_id)
 
         try:
             await self.controller.create_group(leader_id, member_ids)
@@ -285,14 +283,8 @@ class GroupManager:
                 err,
             )
 
-    async def async_unjoin_player(self, player_entity_id: str):
+    async def async_unjoin_player(self, player_id: int, player_entity_id: str):
         """Remove `player_entity_id` from any group."""
-        player_id = self._get_entity_id_to_player_id_map().get(player_entity_id)
-        if not player_id:
-            raise HomeAssistantError(
-                f"The player {player_entity_id} could not be resolved to a HEOS player."
-            )
-
         try:
             await self.controller.create_group(player_id, [])
         except HeosError as err:
@@ -345,6 +337,17 @@ class GroupManager:
             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."""

homeassistant/components/heos/media_player.py

@@ -160,7 +160,11 @@ class HeosMediaPlayer(MediaPlayerEntity):
             async_dispatcher_connect(self.hass, SIGNAL_HEOS_UPDATED, self._heos_updated)
         )
         # Register this player's entity_id so it can be resolved by the group manager
-        self._group_manager.entity_id_map[self._player.player_id] = self.entity_id
+        self.async_on_remove(
+            self._group_manager.register_media_player(
+                self._player.player_id, self.entity_id
+            )
+        )
         async_dispatcher_send(self.hass, SIGNAL_HEOS_PLAYER_ADDED)
 
     @log_command_error("clear playlist")
@@ -171,7 +175,9 @@ class HeosMediaPlayer(MediaPlayerEntity):
     @log_command_error("join_players")
     async def async_join_players(self, group_members: list[str]) -> None:
         """Join `group_members` as a player group with the current player."""
-        await self._group_manager.async_join_players(self.entity_id, group_members)
+        await self._group_manager.async_join_players(
+            self._player.player_id, self.entity_id, group_members
+        )
 
     @log_command_error("pause")
     async def async_media_pause(self) -> None:
@@ -294,7 +300,9 @@ class HeosMediaPlayer(MediaPlayerEntity):
     @log_command_error("unjoin_player")
     async def async_unjoin_player(self) -> None:
         """Remove this player from any group."""
-        await self._group_manager.async_unjoin_player(self.entity_id)
+        await self._group_manager.async_unjoin_player(
+            self._player.player_id, self.entity_id
+        )
 
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect the device when removed."""

tests/components/heos/test_media_player.py

@@ -51,6 +51,7 @@ from homeassistant.const import (
     STATE_UNAVAILABLE,
 )
 from homeassistant.core import HomeAssistant
+from homeassistant.exceptions import HomeAssistantError
 from homeassistant.helpers import device_registry as dr, entity_registry as er
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.setup import async_setup_component
@@ -1051,3 +1052,34 @@ async def test_media_player_unjoin_group(
         blocking=True,
     )
     assert "Failed to ungroup media_player.test_player" in caplog.text
+
+
+async def test_media_player_group_fails_when_entity_removed(
+    hass: HomeAssistant,
+    config_entry,
+    config,
+    controller,
+    entity_registry: er.EntityRegistry,
+    caplog: pytest.LogCaptureFixture,
+) -> None:
+    """Test grouping fails when entity removed."""
+    await setup_platform(hass, config_entry, config)
+
+    # Remove one of the players
+    entity_registry.async_remove("media_player.test_player_2")
+
+    # Attempt to group
+    with pytest.raises(
+        HomeAssistantError,
+        match="The group member media_player.test_player_2 could not be resolved to a HEOS player.",
+    ):
+        await hass.services.async_call(
+            MEDIA_PLAYER_DOMAIN,
+            SERVICE_JOIN,
+            {
+                ATTR_ENTITY_ID: "media_player.test_player",
+                ATTR_GROUP_MEMBERS: ["media_player.test_player_2"],
+            },
+            blocking=True,
+        )
+    controller.create_group.assert_not_called()