Reinitialize hassio discovery flow on config entry removal (#127088)

* Reinitialize hassio discovery flow on config entry removal

* Address review comments

homeassistant/components/hassio/discovery.py

@@ -16,8 +16,9 @@ from homeassistant.const import ATTR_SERVICE, EVENT_HOMEASSISTANT_START
 from homeassistant.core import Event, HomeAssistant, callback
 from homeassistant.data_entry_flow import BaseServiceInfo
 from homeassistant.helpers import discovery_flow
+from homeassistant.helpers.dispatcher import async_dispatcher_connect
 
-from .const import ATTR_ADDON, ATTR_CONFIG, ATTR_DISCOVERY, ATTR_UUID
+from .const import ATTR_ADDON, ATTR_CONFIG, ATTR_DISCOVERY, ATTR_UUID, DOMAIN
 from .handler import HassIO, HassioAPIError
 
 _LOGGER = logging.getLogger(__name__)
@@ -59,6 +60,23 @@ def async_setup_discovery_view(hass: HomeAssistant, hassio: HassIO) -> None:
         EVENT_HOMEASSISTANT_START, _async_discovery_start_handler
     )
 
+    async def _handle_config_entry_removed(
+        entry: config_entries.ConfigEntry,
+    ) -> None:
+        """Handle config entry changes."""
+        for disc_key in entry.discovery_keys[DOMAIN]:
+            if disc_key.version != 1 or not isinstance(key := disc_key.key, str):
+                continue
+            uuid = key
+            _LOGGER.debug("Rediscover addon %s", uuid)
+            await hassio_discovery.async_rediscover(uuid)
+
+    async_dispatcher_connect(
+        hass,
+        config_entries.signal_discovered_config_entry_removed(DOMAIN),
+        _handle_config_entry_removed,
+    )
+
 
 class HassIODiscovery(HomeAssistantView):
     """Hass.io view to handle base part."""
@@ -90,6 +108,15 @@ class HassIODiscovery(HomeAssistantView):
         await self.async_process_del(data)
         return web.Response()
 
+    async def async_rediscover(self, uuid: str) -> None:
+        """Rediscover add-on when config entry is removed."""
+        try:
+            data = await self.hassio.get_discovery_message(uuid)
+        except HassioAPIError as err:
+            _LOGGER.debug("Can't read discovery data: %s", err)
+        else:
+            await self.async_process_new(data)
+
     async def async_process_new(self, data: dict[str, Any]) -> None:
         """Process add discovery entry."""
         service: str = data[ATTR_SERVICE]
@@ -114,6 +141,11 @@ class HassIODiscovery(HomeAssistantView):
             data=HassioServiceInfo(
                 config=config_data, name=addon_info.name, slug=slug, uuid=uuid
             ),
+            discovery_key=discovery_flow.DiscoveryKey(
+                domain=DOMAIN,
+                key=data[ATTR_UUID],
+                version=1,
+            ),
         )
 
     async def async_process_del(self, data: dict[str, Any]) -> None:

tests/components/hassio/test_discovery.py

@@ -13,9 +13,16 @@ from homeassistant.components.hassio.handler import HassioAPIError
 from homeassistant.components.mqtt import DOMAIN as MQTT_DOMAIN
 from homeassistant.const import EVENT_HOMEASSISTANT_START, EVENT_HOMEASSISTANT_STARTED
 from homeassistant.core import HomeAssistant
+from homeassistant.helpers.discovery_flow import DiscoveryKey
 from homeassistant.setup import async_setup_component
 
-from tests.common import MockModule, mock_config_flow, mock_integration, mock_platform
+from tests.common import (
+    MockConfigEntry,
+    MockModule,
+    mock_config_flow,
+    mock_integration,
+    mock_platform,
+)
 from tests.test_util.aiohttp import AiohttpClientMocker
 
 
@@ -218,3 +225,154 @@ async def test_hassio_discovery_webhook(
             uuid="test",
         )
     )
+
+
+@pytest.mark.parametrize(
+    (
+        "entry_domain",
+        "entry_discovery_keys",
+    ),
+    [
+        # Matching discovery key
+        (
+            "mock-domain",
+            {"hassio": (DiscoveryKey(domain="hassio", key="test", version=1),)},
+        ),
+        # Matching discovery key
+        (
+            "mock-domain",
+            {
+                "hassio": (DiscoveryKey(domain="hassio", key="test", version=1),),
+                "other": (DiscoveryKey(domain="other", key="blah", version=1),),
+            },
+        ),
+        # Matching discovery key, other domain
+        # Note: Rediscovery is not currently restricted to the domain of the removed
+        # entry. Such a check can be added if needed.
+        (
+            "comp",
+            {"hassio": (DiscoveryKey(domain="hassio", key="test", version=1),)},
+        ),
+    ],
+)
+@pytest.mark.parametrize(
+    "entry_source",
+    [
+        config_entries.SOURCE_HASSIO,
+        config_entries.SOURCE_IGNORE,
+        config_entries.SOURCE_USER,
+    ],
+)
+async def test_hassio_rediscover(
+    hass: HomeAssistant,
+    aioclient_mock: AiohttpClientMocker,
+    hassio_client: TestClient,
+    addon_installed: AsyncMock,
+    entry_domain: str,
+    entry_discovery_keys: dict[str, tuple[DiscoveryKey, ...]],
+    entry_source: str,
+) -> None:
+    """Test we reinitiate flows when an ignored config entry is removed."""
+
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+    await hass.async_block_till_done()
+
+    entry = MockConfigEntry(
+        domain=entry_domain,
+        discovery_keys=entry_discovery_keys,
+        unique_id="mock-unique-id",
+        state=config_entries.ConfigEntryState.LOADED,
+        source=entry_source,
+    )
+    entry.add_to_hass(hass)
+
+    aioclient_mock.get(
+        "http://127.0.0.1/discovery/test",
+        json={
+            "result": "ok",
+            "data": {
+                "service": "mqtt",
+                "uuid": "test",
+                "addon": "mosquitto",
+                "config": {
+                    "broker": "mock-broker",
+                    "port": 1883,
+                    "username": "mock-user",
+                    "password": "mock-pass",
+                    "protocol": "3.1.1",
+                },
+            },
+        },
+    )
+    aioclient_mock.get(
+        "http://127.0.0.1/discovery", json={"result": "ok", "data": {"discovery": []}}
+    )
+
+    expected_context = {
+        "discovery_key": DiscoveryKey(domain="hassio", key="test", version=1),
+        "source": config_entries.SOURCE_HASSIO,
+    }
+
+    with patch.object(hass.config_entries.flow, "async_init") as mock_init:
+        await hass.config_entries.async_remove(entry.entry_id)
+        await hass.async_block_till_done()
+
+        assert len(mock_init.mock_calls) == 1
+        assert mock_init.mock_calls[0][1][0] == "mqtt"
+        assert mock_init.mock_calls[0][2]["context"] == expected_context
+
+
+@pytest.mark.usefixtures("mock_async_zeroconf")
+@pytest.mark.parametrize(
+    (
+        "entry_domain",
+        "entry_discovery_keys",
+        "entry_source",
+        "entry_unique_id",
+    ),
+    [
+        # Discovery key from other domain
+        (
+            "mock-domain",
+            {"bluetooth": (DiscoveryKey(domain="bluetooth", key="test", version=1),)},
+            config_entries.SOURCE_IGNORE,
+            "mock-unique-id",
+        ),
+        # Discovery key from the future
+        (
+            "mock-domain",
+            {"hassio": (DiscoveryKey(domain="hassio", key="test", version=2),)},
+            config_entries.SOURCE_IGNORE,
+            "mock-unique-id",
+        ),
+    ],
+)
+async def test_hassio_rediscover_no_match(
+    hass: HomeAssistant,
+    hassio_client: TestClient,
+    entry_domain: str,
+    entry_discovery_keys: dict[str, tuple[DiscoveryKey, ...]],
+    entry_source: str,
+    entry_unique_id: str,
+) -> None:
+    """Test we don't reinitiate flows when a non matching config entry is removed."""
+
+    mock_integration(hass, MockModule(entry_domain))
+
+    hass.bus.async_fire(EVENT_HOMEASSISTANT_STARTED)
+    await hass.async_block_till_done()
+
+    entry = MockConfigEntry(
+        domain=entry_domain,
+        discovery_keys=entry_discovery_keys,
+        unique_id=entry_unique_id,
+        state=config_entries.ConfigEntryState.LOADED,
+        source=entry_source,
+    )
+    entry.add_to_hass(hass)
+
+    with patch.object(hass.config_entries.flow, "async_init") as mock_init:
+        await hass.config_entries.async_remove(entry.entry_id)
+        await hass.async_block_till_done()
+
+        assert len(mock_init.mock_calls) == 0