Use the shared zeroconf instance when attempting to create another Zeroconf instance

homeassistant/components/zeroconf/__init__.py

@@ -30,7 +30,7 @@
 from homeassistant.helpers.singleton import singleton
 from homeassistant.loader import async_get_homekit, async_get_zeroconf
 
-from .usage import install_multiple_zeroconf_warning
+from .usage import install_multiple_zeroconf_catcher
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -137,7 +137,7 @@ def setup(hass, config):
         ipv6=zc_config.get(CONF_IPV6, DEFAULT_IPV6),
     )
 
-    install_multiple_zeroconf_warning()
+    install_multiple_zeroconf_catcher(zeroconf)
 
     # Get instance UUID
     uuid = asyncio.run_coroutine_threadsafe(

homeassistant/components/zeroconf/usage.py

@@ -2,13 +2,17 @@
 
 import zeroconf
 
-from homeassistant.helpers.frame import warn_use
+from homeassistant.helpers.frame import report
 
 
-def install_multiple_zeroconf_warning() -> None:
-    """Wrap the Zeroconf class to warn if multiple instances are detected."""
+def install_multiple_zeroconf_catcher(zc) -> None:
+    """Wrap the Zeroconf class to return the shared instance if multiple instances are detected."""
 
-    zeroconf.Zeroconf.__init__ = warn_use(  # type: ignore
-        zeroconf.Zeroconf.__init__,
-        "created another Zeroconf instance. Please use the shared Zeroconf via homeassistant.components.zeroconf.async_get_instance()",
-    )
+    def new_zeroconf_new(self, *k, **kw) -> zeroconf.Zeroconf:  # type: ignore
+        report(
+            "attempted to create another Zeroconf instance. Please use the shared Zeroconf via await homeassistant.components.zeroconf.async_get_instance(hass)",
+            exclude_integrations={"zeroconf"},
+        )
+        return zc
+
+    zeroconf.Zeroconf.__new__ = new_zeroconf_new

homeassistant/helpers/frame.py

@@ -3,7 +3,7 @@
 import functools
 import logging
 from traceback import FrameSummary, extract_stack
-from typing import Any, Callable, Tuple, TypeVar, cast
+from typing import Any, Callable, Optional, Tuple, TypeVar, cast
 
 from homeassistant.exceptions import HomeAssistantError
 
@@ -12,15 +12,24 @@
 CALLABLE_T = TypeVar("CALLABLE_T", bound=Callable)  # pylint: disable=invalid-name
 
 
-def get_integration_frame() -> Tuple[FrameSummary, str, str]:
+def get_integration_frame(
+    exclude_integrations: Optional[set] = None,
+) -> Tuple[FrameSummary, str, str]:
     """Return the frame, integration and integration path of the current stack frame."""
     found_frame = None
+    if not exclude_integrations:
+        exclude_integrations = set()
 
     for frame in reversed(extract_stack()):
         for path in ("custom_components/", "homeassistant/components/"):
             try:
                 index = frame.filename.index(path)
-                found_frame = frame
+                start = index + len(path)
+                end = frame.filename.index("/", start)
+                integration = frame.filename[start:end]
+                if integration not in exclude_integrations:
+                    found_frame = frame
+
                 break
             except ValueError:
                 continue
@@ -31,25 +40,22 @@ def get_integration_frame() -> Tuple[FrameSummary, str, str]:
     if found_frame is None:
         raise MissingIntegrationFrame
 
-    start = index + len(path)
-    end = found_frame.filename.index("/", start)
-
-    integration = found_frame.filename[start:end]
-
     return found_frame, integration, path
 
 
 class MissingIntegrationFrame(HomeAssistantError):
     """Raised when no integration is found in the frame."""
 
 
-def report(what: str) -> None:
+def report(what: str, exclude_integrations: Optional[set] = None) -> None:
     """Report incorrect usage.
 
     Async friendly.
     """
     try:
-        found_frame, integration, path = get_integration_frame()
+        found_frame, integration, path = get_integration_frame(
+            exclude_integrations=exclude_integrations
+        )
     except MissingIntegrationFrame:
         # Did not source from an integration? Hard error.
         raise RuntimeError(f"Detected code that {what}. Please report this issue.")

tests/components/zeroconf/test_init.py

@@ -8,7 +8,7 @@
 from homeassistant.generated import zeroconf as zc_gen
 from homeassistant.setup import async_setup_component
 
-from tests.async_mock import patch
+from tests.async_mock import Mock, patch
 
 NON_UTF8_VALUE = b"ABCDEF\x8a"
 NON_ASCII_KEY = b"non-ascii-key\x8a"
@@ -339,11 +339,48 @@ async def test_get_instance(hass, mock_zeroconf):
     assert len(mock_zeroconf.ha_close.mock_calls) == 1
 
 
-async def test_multiple_zeroconf_instances(hass, mock_zeroconf, caplog):
-    """Test creating multiple zeroconf throws."""
+async def test_multiple_zeroconf_instances(hass, mock_zeroconf):
+    """Test creating multiple zeroconf throws without an integration."""
 
     assert await async_setup_component(hass, zeroconf.DOMAIN, {zeroconf.DOMAIN: {}})
     await hass.async_block_till_done()
 
     with pytest.raises(RuntimeError):
         zeroconf.Zeroconf()
+
+
+async def test_multiple_zeroconf_instances_gives_shared(hass, mock_zeroconf, caplog):
+    """Test creating multiple zeroconf gives the shared instance to an integration."""
+
+    assert await async_setup_component(hass, zeroconf.DOMAIN, {zeroconf.DOMAIN: {}})
+    await hass.async_block_till_done()
+
+    zeroconf_instance = await zeroconf.async_get_instance(hass)
+
+    correct_frame = Mock(
+        filename="/config/custom_components/burncpu/light.py",
+        lineno="23",
+        line="self.light.is_on",
+    )
+    with patch(
+        "homeassistant.helpers.frame.extract_stack",
+        return_value=[
+            Mock(
+                filename="/home/dev/homeassistant/core.py",
+                lineno="23",
+                line="do_something()",
+            ),
+            correct_frame,
+            Mock(
+                filename="/home/dev/homeassistant/components/zeroconf/usage.py",
+                lineno="23",
+                line="self.light.is_on",
+            ),
+            Mock(filename="/home/dev/mdns/lights.py", lineno="2", line="something()",),
+        ],
+    ):
+        assert zeroconf.Zeroconf() == zeroconf_instance
+
+    assert "custom_components/burncpu/light.py" in caplog.text
+    assert "23" in caplog.text
+    assert "self.light.is_on" in caplog.text

tests/helpers/test_frame.py

@@ -36,6 +36,39 @@ async def test_extract_frame_integration(caplog):
     assert found_frame == correct_frame
 
 
+async def test_extract_frame_integration_with_excluded_intergration(caplog):
+    """Test extracting the current frame from integration context."""
+    correct_frame = Mock(
+        filename="/home/dev/homeassistant/components/mdns/light.py",
+        lineno="23",
+        line="self.light.is_on",
+    )
+    with patch(
+        "homeassistant.helpers.frame.extract_stack",
+        return_value=[
+            Mock(
+                filename="/home/dev/homeassistant/core.py",
+                lineno="23",
+                line="do_something()",
+            ),
+            correct_frame,
+            Mock(
+                filename="/home/dev/homeassistant/components/zeroconf/usage.py",
+                lineno="23",
+                line="self.light.is_on",
+            ),
+            Mock(filename="/home/dev/mdns/lights.py", lineno="2", line="something()",),
+        ],
+    ):
+        found_frame, integration, path = frame.get_integration_frame(
+            exclude_integrations={"zeroconf"}
+        )
+
+    assert integration == "mdns"
+    assert path == "homeassistant/components/"
+    assert found_frame == correct_frame
+
+
 async def test_extract_frame_no_integration(caplog):
     """Test extracting the current frame without integration context."""
     with patch(