Update IDs after firmware upgrade in HEOS (#23641)

* Initial work

* Update tests

homeassistant/components/heos/__init__.py

@@ -2,6 +2,7 @@
 import asyncio
 from datetime import timedelta
 import logging
+from typing import Dict
 
 import voluptuous as vol
 
@@ -16,8 +17,8 @@ from homeassistant.util import Throttle
 
 from .config_flow import format_title
 from .const import (
-    COMMAND_RETRY_ATTEMPTS, COMMAND_RETRY_DELAY, DATA_CONTROLLER,
+    COMMAND_RETRY_ATTEMPTS, COMMAND_RETRY_DELAY, DATA_CONTROLLER_MANAGER,
-    DATA_SOURCE_MANAGER, DOMAIN, SIGNAL_HEOS_SOURCES_UPDATED)
+    DATA_SOURCE_MANAGER, DOMAIN, SIGNAL_HEOS_UPDATED)
 
 CONFIG_SCHEMA = vol.Schema({
     DOMAIN: vol.Schema({
@@ -89,11 +90,14 @@ async def async_setup_entry(hass: HomeAssistantType, entry: ConfigEntry):
                       exc_info=isinstance(error, CommandError))
         raise ConfigEntryNotReady
 
+    controller_manager = ControllerManager(hass, controller)
+    await controller_manager.connect_listeners()
+
     source_manager = SourceManager(favorites, inputs)
     source_manager.connect_update(hass, controller)
 
     hass.data[DOMAIN] = {
-        DATA_CONTROLLER: controller,
+        DATA_CONTROLLER_MANAGER: controller_manager,
         DATA_SOURCE_MANAGER: source_manager,
         MEDIA_PLAYER_DOMAIN: players
     }
@@ -104,14 +108,91 @@ async def async_setup_entry(hass: HomeAssistantType, entry: ConfigEntry):
 
 async def async_unload_entry(hass: HomeAssistantType, entry: ConfigEntry):
     """Unload a config entry."""
-    controller = hass.data[DOMAIN][DATA_CONTROLLER]
+    controller_manager = hass.data[DOMAIN][DATA_CONTROLLER_MANAGER]
-    controller.dispatcher.disconnect_all()
+    await controller_manager.disconnect()
-    await controller.disconnect()
     hass.data.pop(DOMAIN)
     return await hass.config_entries.async_forward_entry_unload(
         entry, MEDIA_PLAYER_DOMAIN)
 
 
+class ControllerManager:
+    """Class that manages events of the controller."""
+
+    def __init__(self, hass, controller):
+        """Init the controller manager."""
+        self._hass = hass
+        self._device_registry = None
+        self._entity_registry = None
+        self.controller = controller
+        self._signals = []
+
+    async def connect_listeners(self):
+        """Subscribe to events of interest."""
+        from pyheos import const
+        self._device_registry, self._entity_registry = await asyncio.gather(
+            self._hass.helpers.device_registry.async_get_registry(),
+            self._hass.helpers.entity_registry.async_get_registry())
+        # Handle controller events
+        self._signals.append(self.controller.dispatcher.connect(
+            const.SIGNAL_CONTROLLER_EVENT, self._controller_event))
+        # Handle connection-related events
+        self._signals.append(self.controller.dispatcher.connect(
+            const.SIGNAL_HEOS_EVENT, self._heos_event))
+
+    async def disconnect(self):
+        """Disconnect subscriptions."""
+        for signal_remove in self._signals:
+            signal_remove()
+        self._signals.clear()
+        self.controller.dispatcher.disconnect_all()
+        await self.controller.disconnect()
+
+    async def _controller_event(self, event, data):
+        """Handle controller event."""
+        from pyheos import const
+        if event == const.EVENT_PLAYERS_CHANGED:
+            self.update_ids(data[const.DATA_MAPPED_IDS])
+        # Update players
+        self._hass.helpers.dispatcher.async_dispatcher_send(
+            SIGNAL_HEOS_UPDATED)
+
+    async def _heos_event(self, event):
+        """Handle connection event."""
+        from pyheos import CommandError, const
+        if event == const.EVENT_CONNECTED:
+            try:
+                # Retrieve latest players and refresh status
+                data = await self.controller.load_players()
+                self.update_ids(data[const.DATA_MAPPED_IDS])
+            except (CommandError, asyncio.TimeoutError, ConnectionError) as ex:
+                _LOGGER.error("Unable to refresh players: %s", ex)
+        # Update players
+        self._hass.helpers.dispatcher.async_dispatcher_send(
+            SIGNAL_HEOS_UPDATED)
+
+    def update_ids(self, mapped_ids: Dict[int, int]):
+        """Update the IDs in the device and entity registry."""
+        # mapped_ids contains the mapped IDs (new:old)
+        for new_id, old_id in mapped_ids.items():
+            # update device registry
+            entry = self._device_registry.async_get_device(
+                {(DOMAIN, old_id)}, set())
+            new_identifiers = {(DOMAIN, new_id)}
+            if entry:
+                self._device_registry.async_update_device(
+                    entry.id, new_identifiers=new_identifiers)
+                _LOGGER.debug("Updated device %s identifiers to %s",
+                              entry.id, new_identifiers)
+            # update entity registry
+            entity_id = self._entity_registry.async_get_entity_id(
+                MEDIA_PLAYER_DOMAIN, DOMAIN, str(old_id))
+            if entity_id:
+                self._entity_registry.async_update_entity(
+                    entity_id, new_unique_id=str(new_id))
+                _LOGGER.debug("Updated entity %s unique id to %s",
+                              entity_id, new_id)
+
+
 class SourceManager:
     """Class that manages sources for players."""
 
@@ -195,9 +276,10 @@ class SourceManager:
                                       exc_info=isinstance(error, CommandError))
                         return
 
-        async def update_sources(event, data):
+        async def update_sources(event, data=None):
             if event in (const.EVENT_SOURCES_CHANGED,
-                         const.EVENT_USER_CHANGED):
+                         const.EVENT_USER_CHANGED,
+                         const.EVENT_CONNECTED):
                 sources = await get_sources()
                 # If throttled, it will return None
                 if sources:
@@ -206,7 +288,9 @@ class SourceManager:
                     _LOGGER.debug("Sources updated due to changed event")
                     # Let players know to update
                     hass.helpers.dispatcher.async_dispatcher_send(
-                        SIGNAL_HEOS_SOURCES_UPDATED)
+                        SIGNAL_HEOS_UPDATED)
 
         controller.dispatcher.connect(
             const.SIGNAL_CONTROLLER_EVENT, update_sources)
+        controller.dispatcher.connect(
+            const.SIGNAL_HEOS_EVENT, update_sources)

homeassistant/components/heos/const.py

@@ -2,8 +2,8 @@
 
 COMMAND_RETRY_ATTEMPTS = 2
 COMMAND_RETRY_DELAY = 1
-DATA_CONTROLLER = "controller"
+DATA_CONTROLLER_MANAGER = "controller"
 DATA_SOURCE_MANAGER = "source_manager"
 DATA_DISCOVERED_HOSTS = "heos_discovered_hosts"
 DOMAIN = 'heos'
-SIGNAL_HEOS_SOURCES_UPDATED = "heos_sources_updated"
+SIGNAL_HEOS_UPDATED = "heos_updated"

homeassistant/components/heos/media_player.py

@@ -18,7 +18,7 @@ from homeassistant.helpers.typing import HomeAssistantType
 from homeassistant.util.dt import utcnow
 
 from .const import (
-    DATA_SOURCE_MANAGER, DOMAIN as HEOS_DOMAIN, SIGNAL_HEOS_SOURCES_UPDATED)
+    DATA_SOURCE_MANAGER, DOMAIN as HEOS_DOMAIN, SIGNAL_HEOS_UPDATED)
 
 BASE_SUPPORTED_FEATURES = SUPPORT_VOLUME_MUTE | SUPPORT_VOLUME_SET | \
                           SUPPORT_VOLUME_STEP | SUPPORT_CLEAR_PLAYLIST | \
@@ -81,23 +81,6 @@ class HeosMediaPlayer(MediaPlayerDevice):
             const.CONTROL_PLAY_NEXT: SUPPORT_NEXT_TRACK
         }
 
-    async def _controller_event(self, event, data):
-        """Handle controller event."""
-        from pyheos import const
-        if event == const.EVENT_PLAYERS_CHANGED:
-            await self.async_update_ha_state(True)
-
-    async def _heos_event(self, event):
-        """Handle connection event."""
-        from pyheos import CommandError, const
-        if event == const.EVENT_CONNECTED:
-            try:
-                await self._player.refresh()
-            except (CommandError, asyncio.TimeoutError, ConnectionError) as ex:
-                _LOGGER.error("Unable to refresh player %s: %s",
-                              self._player, ex)
-        await self.async_update_ha_state(True)
-
     async def _player_update(self, player_id, event):
         """Handle player attribute updated."""
         from pyheos import const
@@ -107,7 +90,7 @@ class HeosMediaPlayer(MediaPlayerDevice):
             self._media_position_updated_at = utcnow()
         await self.async_update_ha_state(True)
 
-    async def _sources_updated(self):
+    async def _heos_updated(self):
         """Handle sources changed."""
         await self.async_update_ha_state(True)
 
@@ -118,16 +101,10 @@ class HeosMediaPlayer(MediaPlayerDevice):
         # Update state when attributes of the player change
         self._signals.append(self._player.heos.dispatcher.connect(
             const.SIGNAL_PLAYER_EVENT, self._player_update))
-        # Update state when available players change
+        # Update state when heos changes
-        self._signals.append(self._player.heos.dispatcher.connect(
-            const.SIGNAL_CONTROLLER_EVENT, self._controller_event))
-        # Update state upon connect/disconnects
-        self._signals.append(self._player.heos.dispatcher.connect(
-            const.SIGNAL_HEOS_EVENT, self._heos_event))
-        # Update state when sources change
         self._signals.append(
             self.hass.helpers.dispatcher.async_dispatcher_connect(
-                SIGNAL_HEOS_SOURCES_UPDATED, self._sources_updated))
+                SIGNAL_HEOS_UPDATED, self._heos_updated))
 
     @log_command_error("clear playlist")
     async def async_clear_playlist(self):
@@ -252,7 +229,7 @@ class HeosMediaPlayer(MediaPlayerDevice):
         """Get attributes about the device."""
         return {
             'identifiers': {
-                (DOMAIN, self._player.player_id)
+                (HEOS_DOMAIN, self._player.player_id)
             },
             'name': self._player.name,
             'model': self._player.model,

tests/components/heos/conftest.py

@@ -20,7 +20,7 @@ def config_entry_fixture():
 
 @pytest.fixture(name="controller")
 def controller_fixture(
-        players, favorites, input_sources, playlists, dispatcher):
+        players, favorites, input_sources, playlists, change_data, dispatcher):
     """Create a mock Heos controller fixture."""
     with patch("pyheos.Heos", autospec=True) as mock:
         mock_heos = mock.return_value
@@ -32,6 +32,7 @@ def controller_fixture(
         mock_heos.get_favorites.return_value = favorites
         mock_heos.get_input_sources.return_value = input_sources
         mock_heos.get_playlists.return_value = playlists
+        mock_heos.load_players.return_value = change_data
         mock_heos.is_signed_in = True
         mock_heos.signed_in_username = "user@user.com"
         yield mock_heos
@@ -149,3 +150,23 @@ def playlists_fixture() -> Sequence[HeosSource]:
     playlist.type = const.TYPE_PLAYLIST
     playlist.name = "Awesome Music"
     return [playlist]
+
+
+@pytest.fixture(name="change_data")
+def change_data_fixture() -> Dict:
+    """Create player change data for testing."""
+    return {
+        const.DATA_MAPPED_IDS: {},
+        const.DATA_NEW: []
+    }
+
+
+@pytest.fixture(name="change_data_mapped_ids")
+def change_data_mapped_ids_fixture() -> Dict:
+    """Create player change data for testing."""
+    return {
+        const.DATA_MAPPED_IDS: {
+            101: 1
+        },
+        const.DATA_NEW: []
+    }

tests/components/heos/test_init.py

@@ -1,13 +1,14 @@
 """Tests for the init module."""
 import asyncio
 
-from asynctest import patch
+from asynctest import Mock, patch
 from pyheos import CommandError, const
 import pytest
 
-from homeassistant.components.heos import async_setup_entry, async_unload_entry
+from homeassistant.components.heos import (
+    ControllerManager, async_setup_entry, async_unload_entry)
 from homeassistant.components.heos.const import (
-    DATA_CONTROLLER, DATA_SOURCE_MANAGER, DOMAIN)
+    DATA_CONTROLLER_MANAGER, DATA_SOURCE_MANAGER, DOMAIN)
 from homeassistant.components.media_player.const import (
     DOMAIN as MEDIA_PLAYER_DOMAIN)
 from homeassistant.const import CONF_HOST
@@ -74,7 +75,7 @@ async def test_async_setup_entry_loads_platforms(
         assert controller.get_favorites.call_count == 1
         assert controller.get_input_sources.call_count == 1
         controller.disconnect.assert_not_called()
-    assert hass.data[DOMAIN][DATA_CONTROLLER] == controller
+    assert hass.data[DOMAIN][DATA_CONTROLLER_MANAGER].controller == controller
     assert hass.data[DOMAIN][MEDIA_PLAYER_DOMAIN] == controller.players
     assert hass.data[DOMAIN][DATA_SOURCE_MANAGER].favorites == favorites
     assert hass.data[DOMAIN][DATA_SOURCE_MANAGER].inputs == input_sources
@@ -97,7 +98,7 @@ async def test_async_setup_entry_not_signed_in_loads_platforms(
         assert controller.get_favorites.call_count == 0
         assert controller.get_input_sources.call_count == 1
         controller.disconnect.assert_not_called()
-    assert hass.data[DOMAIN][DATA_CONTROLLER] == controller
+    assert hass.data[DOMAIN][DATA_CONTROLLER_MANAGER].controller == controller
     assert hass.data[DOMAIN][MEDIA_PLAYER_DOMAIN] == controller.players
     assert hass.data[DOMAIN][DATA_SOURCE_MANAGER].favorites == {}
     assert hass.data[DOMAIN][DATA_SOURCE_MANAGER].inputs == input_sources
@@ -139,12 +140,13 @@ async def test_async_setup_entry_player_failure(
 
 async def test_unload_entry(hass, config_entry, controller):
     """Test entries are unloaded correctly."""
-    hass.data[DOMAIN] = {DATA_CONTROLLER: controller}
+    controller_manager = Mock(ControllerManager)
+    hass.data[DOMAIN] = {DATA_CONTROLLER_MANAGER: controller_manager}
     with patch.object(hass.config_entries, 'async_forward_entry_unload',
                       return_value=True) as unload:
         assert await async_unload_entry(hass, config_entry)
         await hass.async_block_till_done()
-        assert controller.disconnect.call_count == 1
+        assert controller_manager.disconnect.call_count == 1
         assert unload.call_count == 1
     assert DOMAIN not in hass.data
 

tests/components/heos/test_media_player.py

@@ -5,7 +5,7 @@ from pyheos import CommandError, const
 
 from homeassistant.components.heos import media_player
 from homeassistant.components.heos.const import (
-    DATA_SOURCE_MANAGER, DOMAIN, SIGNAL_HEOS_SOURCES_UPDATED)
+    DATA_SOURCE_MANAGER, DOMAIN, SIGNAL_HEOS_UPDATED)
 from homeassistant.components.media_player.const import (
     ATTR_INPUT_SOURCE, ATTR_INPUT_SOURCE_LIST, ATTR_MEDIA_ALBUM_NAME,
     ATTR_MEDIA_ARTIST, ATTR_MEDIA_CONTENT_ID, ATTR_MEDIA_CONTENT_TYPE,
@@ -66,7 +66,7 @@ async def test_state_attributes(hass, config_entry, config, controller):
         hass.data[DOMAIN][DATA_SOURCE_MANAGER].source_list
 
 
-async def test_updates_start_from_signals(
+async def test_updates_from_signals(
         hass, config_entry, config, controller, favorites):
     """Tests dispatched signals update player."""
     await setup_platform(hass, config_entry, config)
@@ -102,48 +102,53 @@ async def test_updates_start_from_signals(
     assert state.attributes[ATTR_MEDIA_DURATION] == 360
     assert state.attributes[ATTR_MEDIA_POSITION] == 1
 
-    # Test controller player change updates
-    player.available = False
-    player.heos.dispatcher.send(
-        const.SIGNAL_CONTROLLER_EVENT, const.EVENT_PLAYERS_CHANGED, {})
-    await hass.async_block_till_done()
-    state = hass.states.get('media_player.test_player')
-    assert state.state == STATE_UNAVAILABLE
-
 
 async def test_updates_from_connection_event(
-        hass, config_entry, config, controller, input_sources, caplog):
+        hass, config_entry, config, controller, caplog):
     """Tests player updates from connection event after connection failure."""
-    # Connected
     await setup_platform(hass, config_entry, config)
     player = controller.players[1]
+    event = asyncio.Event()
+
+    async def set_signal():
+        event.set()
+    hass.helpers.dispatcher.async_dispatcher_connect(
+        SIGNAL_HEOS_UPDATED, set_signal)
+
+    # Connected
     player.available = True
     player.heos.dispatcher.send(
         const.SIGNAL_HEOS_EVENT, const.EVENT_CONNECTED)
-    await hass.async_block_till_done()
+    await event.wait()
     state = hass.states.get('media_player.test_player')
     assert state.state == STATE_IDLE
-    assert player.refresh.call_count == 1
+    assert controller.load_players.call_count == 1
-
-    # Connected handles refresh failure
-    player.reset_mock()
-    player.refresh.side_effect = CommandError(None, "Failure", 1)
-    player.heos.dispatcher.send(
-        const.SIGNAL_HEOS_EVENT, const.EVENT_CONNECTED)
-    await hass.async_block_till_done()
-    state = hass.states.get('media_player.test_player')
-    assert player.refresh.call_count == 1
-    assert "Unable to refresh player" in caplog.text
 
     # Disconnected
+    event.clear()
     player.reset_mock()
+    controller.load_players.reset_mock()
     player.available = False
     player.heos.dispatcher.send(
         const.SIGNAL_HEOS_EVENT, const.EVENT_DISCONNECTED)
-    await hass.async_block_till_done()
+    await event.wait()
     state = hass.states.get('media_player.test_player')
     assert state.state == STATE_UNAVAILABLE
-    assert player.refresh.call_count == 0
+    assert controller.load_players.call_count == 0
+
+    # Connected handles refresh failure
+    event.clear()
+    player.reset_mock()
+    controller.load_players.reset_mock()
+    controller.load_players.side_effect = CommandError(None, "Failure", 1)
+    player.available = True
+    player.heos.dispatcher.send(
+        const.SIGNAL_HEOS_EVENT, const.EVENT_CONNECTED)
+    await event.wait()
+    state = hass.states.get('media_player.test_player')
+    assert state.state == STATE_IDLE
+    assert controller.load_players.call_count == 1
+    assert "Unable to refresh players" in caplog.text
 
 
 async def test_updates_from_sources_updated(
@@ -156,7 +161,7 @@ async def test_updates_from_sources_updated(
     async def set_signal():
         event.set()
     hass.helpers.dispatcher.async_dispatcher_connect(
-        SIGNAL_HEOS_SOURCES_UPDATED, set_signal)
+        SIGNAL_HEOS_UPDATED, set_signal)
 
     input_sources.clear()
     player.heos.dispatcher.send(
@@ -168,6 +173,65 @@ async def test_updates_from_sources_updated(
     assert state.attributes[ATTR_INPUT_SOURCE_LIST] == source_list
 
 
+async def test_updates_from_players_changed(
+        hass, config_entry, config, controller, change_data,
+        caplog):
+    """Test player updates from changes to available players."""
+    await setup_platform(hass, config_entry, config)
+    player = controller.players[1]
+    event = asyncio.Event()
+
+    async def set_signal():
+        event.set()
+    hass.helpers.dispatcher.async_dispatcher_connect(
+        SIGNAL_HEOS_UPDATED, set_signal)
+
+    assert hass.states.get('media_player.test_player').state == STATE_IDLE
+    player.state = const.PLAY_STATE_PLAY
+    player.heos.dispatcher.send(
+        const.SIGNAL_CONTROLLER_EVENT, const.EVENT_PLAYERS_CHANGED,
+        change_data)
+    await event.wait()
+    assert hass.states.get('media_player.test_player').state == STATE_PLAYING
+
+
+async def test_updates_from_players_changed_new_ids(
+        hass, config_entry, config, controller, change_data_mapped_ids,
+        caplog):
+    """Test player updates from changes to available players."""
+    await setup_platform(hass, config_entry, config)
+    device_registry = await hass.helpers.device_registry.async_get_registry()
+    entity_registry = await hass.helpers.entity_registry.async_get_registry()
+    player = controller.players[1]
+    event = asyncio.Event()
+
+    # Assert device registry matches current id
+    assert device_registry.async_get_device(
+        {(DOMAIN, 1)}, [])
+    # Assert entity registry matches current id
+    assert entity_registry.async_get_entity_id(
+        MEDIA_PLAYER_DOMAIN, DOMAIN, '1') == "media_player.test_player"
+
+    # Trigger update
+    async def set_signal():
+        event.set()
+    hass.helpers.dispatcher.async_dispatcher_connect(
+        SIGNAL_HEOS_UPDATED, set_signal)
+    player.heos.dispatcher.send(
+        const.SIGNAL_CONTROLLER_EVENT, const.EVENT_PLAYERS_CHANGED,
+        change_data_mapped_ids)
+    await event.wait()
+
+    # Assert device registry identifiers were updated
+    assert len(device_registry.devices) == 1
+    assert device_registry.async_get_device(
+        {(DOMAIN, 101)}, [])
+    # Assert entity registry unique id was updated
+    assert len(entity_registry.entities) == 1
+    assert entity_registry.async_get_entity_id(
+        MEDIA_PLAYER_DOMAIN, DOMAIN, '101') == "media_player.test_player"
+
+
 async def test_updates_from_user_changed(
         hass, config_entry, config, controller):
     """Tests player updates from changes in user."""
@@ -178,7 +242,7 @@ async def test_updates_from_user_changed(
     async def set_signal():
         event.set()
     hass.helpers.dispatcher.async_dispatcher_connect(
-        SIGNAL_HEOS_SOURCES_UPDATED, set_signal)
+        SIGNAL_HEOS_UPDATED, set_signal)
 
     controller.is_signed_in = False
     controller.signed_in_username = None