Add zeroconf/homekit/ssdp discovery support for custom components (#38466)

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/ssdp/__init__.py

@@ -8,8 +8,8 @@ from defusedxml import ElementTree
 from netdisco import ssdp, util
 
 from homeassistant.const import EVENT_HOMEASSISTANT_STARTED
-from homeassistant.generated.ssdp import SSDP
 from homeassistant.helpers.event import async_track_time_interval
+from homeassistant.loader import async_get_ssdp
 
 DOMAIN = "ssdp"
 SCAN_INTERVAL = timedelta(seconds=60)
@@ -35,7 +35,7 @@ async def async_setup(hass, config):
     """Set up the SSDP integration."""
 
     async def initialize(_):
-        scanner = Scanner(hass)
+        scanner = Scanner(hass, await async_get_ssdp(hass))
         await scanner.async_scan(None)
         async_track_time_interval(hass, scanner.async_scan, SCAN_INTERVAL)
 
@@ -47,10 +47,11 @@ async def async_setup(hass, config):
 class Scanner:
     """Class to manage SSDP scanning."""
 
-    def __init__(self, hass):
+    def __init__(self, hass, integration_matchers):
         """Initialize class."""
         self.hass = hass
         self.seen = set()
+        self._integration_matchers = integration_matchers
         self._description_cache = {}
 
     async def async_scan(self, _):
@@ -121,7 +122,7 @@ class Scanner:
             info.update(await info_req)
 
         domains = set()
-        for domain, matchers in SSDP.items():
+        for domain, matchers in self._integration_matchers.items():
             for matcher in matchers:
                 if all(info.get(k) == v for (k, v) in matcher.items()):
                     domains.add(domain)

homeassistant/components/zeroconf/__init__.py

@@ -25,10 +25,10 @@ from homeassistant.const import (
     EVENT_HOMEASSISTANT_STOP,
     __version__,
 )
-from homeassistant.generated.zeroconf import HOMEKIT, ZEROCONF
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.network import NoURLAvailableError, get_url
 from homeassistant.helpers.singleton import singleton
+from homeassistant.loader import async_get_homekit, async_get_zeroconf
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -197,8 +197,14 @@ def setup(hass, config):
 
     hass.bus.listen_once(EVENT_HOMEASSISTANT_START, zeroconf_hass_start)
 
+    zeroconf_types = {}
+    homekit_models = {}
+
     def service_update(zeroconf, service_type, name, state_change):
         """Service state changed."""
+        nonlocal zeroconf_types
+        nonlocal homekit_models
+
         if state_change != ServiceStateChange.Added:
             return
 
@@ -219,7 +225,7 @@ def setup(hass, config):
 
         # If we can handle it as a HomeKit discovery, we do that here.
         if service_type == HOMEKIT_TYPE:
-            discovery_was_forwarded = handle_homekit(hass, info)
+            discovery_was_forwarded = handle_homekit(hass, homekit_models, info)
             # Continue on here as homekit_controller
             # still needs to get updates on devices
             # so it can see when the 'c#' field is updated.
@@ -241,20 +247,25 @@ def setup(hass, config):
                     # likely bad homekit data
                     return
 
-        for domain in ZEROCONF[service_type]:
+        for domain in zeroconf_types[service_type]:
             hass.add_job(
                 hass.config_entries.flow.async_init(
                     domain, context={"source": DOMAIN}, data=info
                 )
             )
 
-    types = list(ZEROCONF)
+    async def zeroconf_hass_started(_event):
-
-    if HOMEKIT_TYPE not in ZEROCONF:
-        types.append(HOMEKIT_TYPE)
-
-    def zeroconf_hass_started(_event):
         """Start the service browser."""
+        nonlocal zeroconf_types
+        nonlocal homekit_models
+
+        zeroconf_types = await async_get_zeroconf(hass)
+        homekit_models = await async_get_homekit(hass)
+
+        types = list(zeroconf_types)
+
+        if HOMEKIT_TYPE not in zeroconf_types:
+            types.append(HOMEKIT_TYPE)
 
         _LOGGER.debug("Starting Zeroconf browser")
         HaServiceBrowser(zeroconf, types, handlers=[service_update])
@@ -264,7 +275,7 @@ def setup(hass, config):
     return True
 
 
-def handle_homekit(hass, info) -> bool:
+def handle_homekit(hass, homekit_models, info) -> bool:
     """Handle a HomeKit discovery.
 
     Return if discovery was forwarded.
@@ -280,7 +291,7 @@ def handle_homekit(hass, info) -> bool:
     if model is None:
         return False
 
-    for test_model in HOMEKIT:
+    for test_model in homekit_models:
         if (
             model != test_model
             and not model.startswith(f"{test_model} ")
@@ -290,7 +301,7 @@ def handle_homekit(hass, info) -> bool:
 
         hass.add_job(
             hass.config_entries.flow.async_init(
-                HOMEKIT[test_model], context={"source": "homekit"}, data=info
+                homekit_models[test_model], context={"source": "homekit"}, data=info
             )
         )
         return True

homeassistant/loader.py

@@ -25,6 +25,9 @@ from typing import (
     cast,
 )
 
+from homeassistant.generated.ssdp import SSDP
+from homeassistant.generated.zeroconf import HOMEKIT, ZEROCONF
+
 # Typing imports that create a circular dependency
 if TYPE_CHECKING:
     from homeassistant.core import HomeAssistant
@@ -142,6 +145,56 @@ async def async_get_config_flows(hass: "HomeAssistant") -> Set[str]:
     return flows
 
 
+async def async_get_zeroconf(hass: "HomeAssistant") -> Dict[str, List]:
+    """Return cached list of zeroconf types."""
+    zeroconf: Dict[str, List] = ZEROCONF.copy()
+
+    integrations = await async_get_custom_components(hass)
+    for integration in integrations.values():
+        if not integration.zeroconf:
+            continue
+        for typ in integration.zeroconf:
+            zeroconf.setdefault(typ, [])
+            if integration.domain not in zeroconf[typ]:
+                zeroconf[typ].append(integration.domain)
+
+    return zeroconf
+
+
+async def async_get_homekit(hass: "HomeAssistant") -> Dict[str, str]:
+    """Return cached list of homekit models."""
+
+    homekit: Dict[str, str] = HOMEKIT.copy()
+
+    integrations = await async_get_custom_components(hass)
+    for integration in integrations.values():
+        if (
+            not integration.homekit
+            or "models" not in integration.homekit
+            or not integration.homekit["models"]
+        ):
+            continue
+        for model in integration.homekit["models"]:
+            homekit[model] = integration.domain
+
+    return homekit
+
+
+async def async_get_ssdp(hass: "HomeAssistant") -> Dict[str, List]:
+    """Return cached list of ssdp mappings."""
+
+    ssdp: Dict[str, List] = SSDP.copy()
+
+    integrations = await async_get_custom_components(hass)
+    for integration in integrations.values():
+        if not integration.ssdp:
+            continue
+
+        ssdp[integration.domain] = integration.ssdp
+
+    return ssdp
+
+
 class Integration:
     """An integration in Home Assistant."""
 
@@ -258,6 +311,21 @@ class Integration:
         """Return Integration Quality Scale."""
         return cast(str, self.manifest.get("quality_scale"))
 
+    @property
+    def ssdp(self) -> Optional[list]:
+        """Return Integration SSDP entries."""
+        return cast(List[dict], self.manifest.get("ssdp"))
+
+    @property
+    def zeroconf(self) -> Optional[list]:
+        """Return Integration zeroconf entries."""
+        return cast(List[str], self.manifest.get("zeroconf"))
+
+    @property
+    def homekit(self) -> Optional[dict]:
+        """Return Integration homekit entries."""
+        return cast(Dict[str, List], self.manifest.get("homekit"))
+
     @property
     def is_built_in(self) -> bool:
         """Test if package is a built-in integration."""

tests/components/ssdp/test_init.py

@@ -6,18 +6,17 @@ import aiohttp
 import pytest
 
 from homeassistant.components import ssdp
-from homeassistant.generated import ssdp as gn_ssdp
 
 from tests.common import mock_coro
 
 
 async def test_scan_match_st(hass):
     """Test matching based on ST."""
-    scanner = ssdp.Scanner(hass)
+    scanner = ssdp.Scanner(hass, {"mock-domain": [{"st": "mock-st"}]})
 
     with patch(
         "netdisco.ssdp.scan", return_value=[Mock(st="mock-st", location=None)]
-    ), patch.dict(gn_ssdp.SSDP, {"mock-domain": [{"st": "mock-st"}]}), patch.object(
+    ), patch.object(
         hass.config_entries.flow, "async_init", return_value=mock_coro()
     ) as mock_init:
         await scanner.async_scan(None)
@@ -42,12 +41,12 @@ async def test_scan_match_upnp_devicedesc(hass, aioclient_mock, key):
 </root>
     """,
     )
-    scanner = ssdp.Scanner(hass)
+    scanner = ssdp.Scanner(hass, {"mock-domain": [{key: "Paulus"}]})
 
     with patch(
         "netdisco.ssdp.scan",
         return_value=[Mock(st="mock-st", location="http://1.1.1.1")],
-    ), patch.dict(gn_ssdp.SSDP, {"mock-domain": [{key: "Paulus"}]}), patch.object(
+    ), patch.object(
         hass.config_entries.flow, "async_init", return_value=mock_coro()
     ) as mock_init:
         await scanner.async_scan(None)
@@ -69,13 +68,8 @@ async def test_scan_not_all_present(hass, aioclient_mock):
 </root>
     """,
     )
-    scanner = ssdp.Scanner(hass)
+    scanner = ssdp.Scanner(
-
+        hass,
-    with patch(
-        "netdisco.ssdp.scan",
-        return_value=[Mock(st="mock-st", location="http://1.1.1.1")],
-    ), patch.dict(
-        gn_ssdp.SSDP,
         {
             "mock-domain": [
                 {
@@ -84,6 +78,11 @@ async def test_scan_not_all_present(hass, aioclient_mock):
                 }
             ]
         },
+    )
+
+    with patch(
+        "netdisco.ssdp.scan",
+        return_value=[Mock(st="mock-st", location="http://1.1.1.1")],
     ), patch.object(
         hass.config_entries.flow, "async_init", return_value=mock_coro()
     ) as mock_init:
@@ -105,13 +104,8 @@ async def test_scan_not_all_match(hass, aioclient_mock):
 </root>
     """,
     )
-    scanner = ssdp.Scanner(hass)
+    scanner = ssdp.Scanner(
-
+        hass,
-    with patch(
-        "netdisco.ssdp.scan",
-        return_value=[Mock(st="mock-st", location="http://1.1.1.1")],
-    ), patch.dict(
-        gn_ssdp.SSDP,
         {
             "mock-domain": [
                 {
@@ -120,6 +114,11 @@ async def test_scan_not_all_match(hass, aioclient_mock):
                 }
             ]
         },
+    )
+
+    with patch(
+        "netdisco.ssdp.scan",
+        return_value=[Mock(st="mock-st", location="http://1.1.1.1")],
     ), patch.object(
         hass.config_entries.flow, "async_init", return_value=mock_coro()
     ) as mock_init:
@@ -132,7 +131,7 @@ async def test_scan_not_all_match(hass, aioclient_mock):
 async def test_scan_description_fetch_fail(hass, aioclient_mock, exc):
     """Test failing to fetch description."""
     aioclient_mock.get("http://1.1.1.1", exc=exc)
-    scanner = ssdp.Scanner(hass)
+    scanner = ssdp.Scanner(hass, {})
 
     with patch(
         "netdisco.ssdp.scan",
@@ -149,7 +148,7 @@ async def test_scan_description_parse_fail(hass, aioclient_mock):
 <root>INVALIDXML
     """,
     )
-    scanner = ssdp.Scanner(hass)
+    scanner = ssdp.Scanner(hass, {})
 
     with patch(
         "netdisco.ssdp.scan",

tests/test_loader.py

@@ -168,10 +168,36 @@ def test_integration_properties(hass):
             "domain": "hue",
             "dependencies": ["test-dep"],
             "requirements": ["test-req==1.0.0"],
+            "zeroconf": ["_hue._tcp.local."],
+            "homekit": {"models": ["BSB002"]},
+            "ssdp": [
+                {
+                    "manufacturer": "Royal Philips Electronics",
+                    "modelName": "Philips hue bridge 2012",
+                },
+                {
+                    "manufacturer": "Royal Philips Electronics",
+                    "modelName": "Philips hue bridge 2015",
+                },
+                {"manufacturer": "Signify", "modelName": "Philips hue bridge 2015"},
+            ],
         },
     )
     assert integration.name == "Philips Hue"
     assert integration.domain == "hue"
+    assert integration.homekit == {"models": ["BSB002"]}
+    assert integration.zeroconf == ["_hue._tcp.local."]
+    assert integration.ssdp == [
+        {
+            "manufacturer": "Royal Philips Electronics",
+            "modelName": "Philips hue bridge 2012",
+        },
+        {
+            "manufacturer": "Royal Philips Electronics",
+            "modelName": "Philips hue bridge 2015",
+        },
+        {"manufacturer": "Signify", "modelName": "Philips hue bridge 2015"},
+    ]
     assert integration.dependencies == ["test-dep"]
     assert integration.requirements == ["test-req==1.0.0"]
     assert integration.is_built_in is True
@@ -188,6 +214,9 @@ def test_integration_properties(hass):
         },
     )
     assert integration.is_built_in is False
+    assert integration.homekit is None
+    assert integration.zeroconf is None
+    assert integration.ssdp is None
 
 
 async def test_integrations_only_once(hass):
@@ -217,6 +246,9 @@ def _get_test_integration(hass, name, config_flow):
             "config_flow": config_flow,
             "dependencies": [],
             "requirements": [],
+            "zeroconf": [f"_{name}._tcp.local."],
+            "homekit": {"models": [name]},
+            "ssdp": [{"manufacturer": name, "modelName": name}],
         },
     )
 
@@ -254,6 +286,51 @@ async def test_get_config_flows(hass):
         assert "test_1" not in flows
 
 
+async def test_get_zeroconf(hass):
+    """Verify that custom components with zeroconf are found."""
+    test_1_integration = _get_test_integration(hass, "test_1", True)
+    test_2_integration = _get_test_integration(hass, "test_2", True)
+
+    with patch("homeassistant.loader.async_get_custom_components") as mock_get:
+        mock_get.return_value = {
+            "test_1": test_1_integration,
+            "test_2": test_2_integration,
+        }
+        zeroconf = await loader.async_get_zeroconf(hass)
+        assert zeroconf["_test_1._tcp.local."] == ["test_1"]
+        assert zeroconf["_test_2._tcp.local."] == ["test_2"]
+
+
+async def test_get_homekit(hass):
+    """Verify that custom components with homekit are found."""
+    test_1_integration = _get_test_integration(hass, "test_1", True)
+    test_2_integration = _get_test_integration(hass, "test_2", True)
+
+    with patch("homeassistant.loader.async_get_custom_components") as mock_get:
+        mock_get.return_value = {
+            "test_1": test_1_integration,
+            "test_2": test_2_integration,
+        }
+        homekit = await loader.async_get_homekit(hass)
+        assert homekit["test_1"] == "test_1"
+        assert homekit["test_2"] == "test_2"
+
+
+async def test_get_ssdp(hass):
+    """Verify that custom components with ssdp are found."""
+    test_1_integration = _get_test_integration(hass, "test_1", True)
+    test_2_integration = _get_test_integration(hass, "test_2", True)
+
+    with patch("homeassistant.loader.async_get_custom_components") as mock_get:
+        mock_get.return_value = {
+            "test_1": test_1_integration,
+            "test_2": test_2_integration,
+        }
+        ssdp = await loader.async_get_ssdp(hass)
+        assert ssdp["test_1"] == [{"manufacturer": "test_1", "modelName": "test_1"}]
+        assert ssdp["test_2"] == [{"manufacturer": "test_2", "modelName": "test_2"}]
+
+
 async def test_get_custom_components_safe_mode(hass):
     """Test that we get empty custom components in safe mode."""
     hass.config.safe_mode = True