Improve cleanup of Google Cast entities (#66801)

homeassistant/components/cast/media_player.py

@@ -212,6 +212,10 @@ class CastDevice(MediaPlayerEntity):
     async def async_will_remove_from_hass(self) -> None:
         """Disconnect Chromecast object when removed."""
         await self._async_disconnect()
+        if self._cast_info.uuid is not None:
+            # Remove the entity from the added casts so that it can dynamically
+            # be re-added again.
+            self.hass.data[ADDED_CAST_DEVICES_KEY].remove(self._cast_info.uuid)
         if self._add_remove_handler:
             self._add_remove_handler()
             self._add_remove_handler = None
@@ -253,21 +257,16 @@ class CastDevice(MediaPlayerEntity):
 
     async def _async_disconnect(self):
         """Disconnect Chromecast object if it is set."""
-        if self._chromecast is None:
-            # Can't disconnect if not connected.
-            return
-        _LOGGER.debug(
-            "[%s %s] Disconnecting from chromecast socket",
-            self.entity_id,
-            self._cast_info.friendly_name,
-        )
+        if self._chromecast is not None:
+            _LOGGER.debug(
+                "[%s %s] Disconnecting from chromecast socket",
+                self.entity_id,
+                self._cast_info.friendly_name,
+            )
+            await self.hass.async_add_executor_job(self._chromecast.disconnect)
+
         self._attr_available = False
-        self.async_write_ha_state()
-
-        await self.hass.async_add_executor_job(self._chromecast.disconnect)
-
         self._invalidate()
-
         self.async_write_ha_state()
 
     def _invalidate(self):
@@ -904,16 +903,13 @@ class DynamicCastGroup:
 
     async def _async_disconnect(self):
         """Disconnect Chromecast object if it is set."""
-        if self._chromecast is None:
-            # Can't disconnect if not connected.
-            return
-        _LOGGER.debug(
-            "[%s %s] Disconnecting from chromecast socket",
-            "Dynamic group",
-            self._cast_info.friendly_name,
-        )
-
-        await self.hass.async_add_executor_job(self._chromecast.disconnect)
+        if self._chromecast is not None:
+            _LOGGER.debug(
+                "[%s %s] Disconnecting from chromecast socket",
+                "Dynamic group",
+                self._cast_info.friendly_name,
+            )
+            await self.hass.async_add_executor_job(self._chromecast.disconnect)
 
         self._invalidate()
 

tests/components/cast/test_media_player.py

@@ -38,7 +38,7 @@ from homeassistant.const import (
     EVENT_HOMEASSISTANT_STOP,
 )
 from homeassistant.core import HomeAssistant
-from homeassistant.helpers import entity_registry as er, network
+from homeassistant.helpers import device_registry as dr, entity_registry as er, network
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.setup import async_setup_component
 
@@ -606,6 +606,45 @@ async def test_entity_availability(hass: HomeAssistant):
     assert state.state == "unavailable"
 
 
+@pytest.mark.parametrize("port,entry_type", ((8009, None),))
+async def test_device_registry(hass: HomeAssistant, port, entry_type):
+    """Test device registry integration."""
+    entity_id = "media_player.speaker"
+    reg = er.async_get(hass)
+    dev_reg = dr.async_get(hass)
+
+    info = get_fake_chromecast_info(port=port)
+
+    chromecast, _ = await async_setup_media_player_cast(hass, info)
+    chromecast.cast_type = pychromecast.const.CAST_TYPE_CHROMECAST
+    _, conn_status_cb, _ = get_status_callbacks(chromecast)
+    cast_entry = hass.config_entries.async_entries("cast")[0]
+
+    connection_status = MagicMock()
+    connection_status.status = "CONNECTED"
+    conn_status_cb(connection_status)
+    await hass.async_block_till_done()
+
+    state = hass.states.get(entity_id)
+    assert state is not None
+    assert state.name == "Speaker"
+    assert state.state == "off"
+    assert entity_id == reg.async_get_entity_id("media_player", "cast", str(info.uuid))
+    entity_entry = reg.async_get(entity_id)
+    assert entity_entry.device_id is not None
+    device_entry = dev_reg.async_get(entity_entry.device_id)
+    assert device_entry.entry_type == entry_type
+
+    # Check that the chromecast object is torn down when the device is removed
+    chromecast.disconnect.assert_not_called()
+    dev_reg.async_update_device(
+        device_entry.id, remove_config_entry_id=cast_entry.entry_id
+    )
+    await hass.async_block_till_done()
+    await hass.async_block_till_done()
+    chromecast.disconnect.assert_called_once()
+
+
 async def test_entity_cast_status(hass: HomeAssistant):
     """Test handling of cast status."""
     entity_id = "media_player.speaker"