From 0daf20c0cce2a60e98952b219f32879d23afaf00 Mon Sep 17 00:00:00 2001
From: Jc2k <john.carr@unrouted.co.uk>
Date: Fri, 11 Feb 2022 19:26:35 +0000
Subject: [PATCH] Prepare for new aiohomekit lifecycle API (#66340)

---
 .strict-typing                                |  1 +
 .../components/homekit_controller/__init__.py | 14 +++----
 .../homekit_controller/config_flow.py         |  6 +--
 .../homekit_controller/manifest.json          |  2 +-
 .../components/homekit_controller/utils.py    | 42 +++++++++++++++++++
 mypy.ini                                      | 11 +++++
 requirements_all.txt                          |  2 +-
 requirements_test_all.txt                     |  2 +-
 tests/components/homekit_controller/common.py |  4 +-
 .../components/homekit_controller/conftest.py |  5 ++-
 .../homekit_controller/test_init.py           | 18 ++++----
 11 files changed, 79 insertions(+), 28 deletions(-)
 create mode 100644 homeassistant/components/homekit_controller/utils.py

diff --git a/.strict-typing b/.strict-typing
index efdccdb1a43..a206f562c17 100644
--- a/.strict-typing
+++ b/.strict-typing
@@ -94,6 +94,7 @@ homeassistant.components.homekit_controller.const
 homeassistant.components.homekit_controller.lock
 homeassistant.components.homekit_controller.select
 homeassistant.components.homekit_controller.storage
+homeassistant.components.homekit_controller.utils
 homeassistant.components.homewizard.*
 homeassistant.components.http.*
 homeassistant.components.huawei_lte.*
diff --git a/homeassistant/components/homekit_controller/__init__.py b/homeassistant/components/homekit_controller/__init__.py
index eeca98167d0..e4f8f715bc8 100644
--- a/homeassistant/components/homekit_controller/__init__.py
+++ b/homeassistant/components/homekit_controller/__init__.py
@@ -14,7 +14,6 @@ from aiohomekit.model.characteristics import (
 )
 from aiohomekit.model.services import Service, ServicesTypes
 
-from homeassistant.components import zeroconf
 from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
 from homeassistant.core import Event, HomeAssistant
@@ -24,8 +23,9 @@ from homeassistant.helpers.typing import ConfigType
 
 from .config_flow import normalize_hkid
 from .connection import HKDevice, valid_serial_number
-from .const import CONTROLLER, ENTITY_MAP, KNOWN_DEVICES, TRIGGERS
+from .const import ENTITY_MAP, KNOWN_DEVICES, TRIGGERS
 from .storage import EntityMapStorage
+from .utils import async_get_controller
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -208,10 +208,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     map_storage = hass.data[ENTITY_MAP] = EntityMapStorage(hass)
     await map_storage.async_initialize()
 
-    async_zeroconf_instance = await zeroconf.async_get_async_instance(hass)
-    hass.data[CONTROLLER] = aiohomekit.Controller(
-        async_zeroconf_instance=async_zeroconf_instance
-    )
+    await async_get_controller(hass)
+
     hass.data[KNOWN_DEVICES] = {}
     hass.data[TRIGGERS] = {}
 
@@ -246,10 +244,10 @@ async def async_remove_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
     # Remove cached type data from .storage/homekit_controller-entity-map
     hass.data[ENTITY_MAP].async_delete_map(hkid)
 
+    controller = await async_get_controller(hass)
+
     # Remove the pairing on the device, making the device discoverable again.
     # Don't reuse any objects in hass.data as they are already unloaded
-    async_zeroconf_instance = await zeroconf.async_get_async_instance(hass)
-    controller = aiohomekit.Controller(async_zeroconf_instance=async_zeroconf_instance)
     controller.load_pairing(hkid, dict(entry.data))
     try:
         await controller.remove_pairing(hkid)
diff --git a/homeassistant/components/homekit_controller/config_flow.py b/homeassistant/components/homekit_controller/config_flow.py
index 54d265a5f4a..9053c79b939 100644
--- a/homeassistant/components/homekit_controller/config_flow.py
+++ b/homeassistant/components/homekit_controller/config_flow.py
@@ -20,6 +20,7 @@ from homeassistant.helpers.device_registry import (
 )
 
 from .const import DOMAIN, KNOWN_DEVICES
+from .utils import async_get_controller
 
 HOMEKIT_DIR = ".homekit"
 HOMEKIT_BRIDGE_DOMAIN = "homekit"
@@ -104,10 +105,7 @@ class HomekitControllerFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
 
     async def _async_setup_controller(self):
         """Create the controller."""
-        async_zeroconf_instance = await zeroconf.async_get_async_instance(self.hass)
-        self.controller = aiohomekit.Controller(
-            async_zeroconf_instance=async_zeroconf_instance
-        )
+        self.controller = await async_get_controller(self.hass)
 
     async def async_step_user(self, user_input=None):
         """Handle a flow start."""
diff --git a/homeassistant/components/homekit_controller/manifest.json b/homeassistant/components/homekit_controller/manifest.json
index 8cc6ce2b575..ce7ae876e03 100644
--- a/homeassistant/components/homekit_controller/manifest.json
+++ b/homeassistant/components/homekit_controller/manifest.json
@@ -3,7 +3,7 @@
   "name": "HomeKit Controller",
   "config_flow": true,
   "documentation": "https://www.home-assistant.io/integrations/homekit_controller",
-  "requirements": ["aiohomekit==0.7.5"],
+  "requirements": ["aiohomekit==0.7.7"],
   "zeroconf": ["_hap._tcp.local."],
   "after_dependencies": ["zeroconf"],
   "codeowners": ["@Jc2k", "@bdraco"],
diff --git a/homeassistant/components/homekit_controller/utils.py b/homeassistant/components/homekit_controller/utils.py
new file mode 100644
index 00000000000..6831c3cee4a
--- /dev/null
+++ b/homeassistant/components/homekit_controller/utils.py
@@ -0,0 +1,42 @@
+"""Helper functions for the homekit_controller component."""
+from typing import cast
+
+from aiohomekit import Controller
+
+from homeassistant.components import zeroconf
+from homeassistant.const import EVENT_HOMEASSISTANT_STOP
+from homeassistant.core import Event, HomeAssistant
+
+from .const import CONTROLLER
+
+
+async def async_get_controller(hass: HomeAssistant) -> Controller:
+    """Get or create an aiohomekit Controller instance."""
+    if existing := hass.data.get(CONTROLLER):
+        return cast(Controller, existing)
+
+    async_zeroconf_instance = await zeroconf.async_get_async_instance(hass)
+
+    # In theory another call to async_get_controller could have run while we were
+    # trying to get the zeroconf instance. So we check again to make sure we
+    # don't leak a Controller instance here.
+    if existing := hass.data.get(CONTROLLER):
+        return cast(Controller, existing)
+
+    controller = Controller(async_zeroconf_instance=async_zeroconf_instance)
+
+    hass.data[CONTROLLER] = controller
+
+    async def _async_stop_homekit_controller(event: Event) -> None:
+        # Pop first so that in theory another controller /could/ start
+        # While this one was shutting down
+        hass.data.pop(CONTROLLER, None)
+        await controller.async_stop()
+
+    # Right now _async_stop_homekit_controller is only called on HA exiting
+    # So we don't have to worry about leaking a callback here.
+    hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, _async_stop_homekit_controller)
+
+    await controller.async_start()
+
+    return controller
diff --git a/mypy.ini b/mypy.ini
index d4db04e8e82..bbab7d20f80 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -851,6 +851,17 @@ no_implicit_optional = true
 warn_return_any = true
 warn_unreachable = true
 
+[mypy-homeassistant.components.homekit_controller.utils]
+check_untyped_defs = true
+disallow_incomplete_defs = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_decorators = true
+disallow_untyped_defs = true
+no_implicit_optional = true
+warn_return_any = true
+warn_unreachable = true
+
 [mypy-homeassistant.components.homewizard.*]
 check_untyped_defs = true
 disallow_incomplete_defs = true
diff --git a/requirements_all.txt b/requirements_all.txt
index 7d05f5eca7b..6d373252f81 100644
--- a/requirements_all.txt
+++ b/requirements_all.txt
@@ -184,7 +184,7 @@ aioguardian==2021.11.0
 aioharmony==0.2.9
 
 # homeassistant.components.homekit_controller
-aiohomekit==0.7.5
+aiohomekit==0.7.7
 
 # homeassistant.components.emulated_hue
 # homeassistant.components.http
diff --git a/requirements_test_all.txt b/requirements_test_all.txt
index 983b50d9e7f..c6ff0311cca 100644
--- a/requirements_test_all.txt
+++ b/requirements_test_all.txt
@@ -134,7 +134,7 @@ aioguardian==2021.11.0
 aioharmony==0.2.9
 
 # homeassistant.components.homekit_controller
-aiohomekit==0.7.5
+aiohomekit==0.7.7
 
 # homeassistant.components.emulated_hue
 # homeassistant.components.http
diff --git a/tests/components/homekit_controller/common.py b/tests/components/homekit_controller/common.py
index 4fabe504d4b..7fae69ee01b 100644
--- a/tests/components/homekit_controller/common.py
+++ b/tests/components/homekit_controller/common.py
@@ -174,7 +174,9 @@ async def setup_platform(hass):
     """Load the platform but with a fake Controller API."""
     config = {"discovery": {}}
 
-    with mock.patch("aiohomekit.Controller") as controller:
+    with mock.patch(
+        "homeassistant.components.homekit_controller.utils.Controller"
+    ) as controller:
         fake_controller = controller.return_value = FakeController()
         await async_setup_component(hass, DOMAIN, config)
 
diff --git a/tests/components/homekit_controller/conftest.py b/tests/components/homekit_controller/conftest.py
index 46b8a5de3e7..81688f88a4b 100644
--- a/tests/components/homekit_controller/conftest.py
+++ b/tests/components/homekit_controller/conftest.py
@@ -27,7 +27,10 @@ def utcnow(request):
 def controller(hass):
     """Replace aiohomekit.Controller with an instance of aiohomekit.testing.FakeController."""
     instance = FakeController()
-    with unittest.mock.patch("aiohomekit.Controller", return_value=instance):
+    with unittest.mock.patch(
+        "homeassistant.components.homekit_controller.utils.Controller",
+        return_value=instance,
+    ):
         yield instance
 
 
diff --git a/tests/components/homekit_controller/test_init.py b/tests/components/homekit_controller/test_init.py
index d1b133468d5..03694e7186a 100644
--- a/tests/components/homekit_controller/test_init.py
+++ b/tests/components/homekit_controller/test_init.py
@@ -4,7 +4,6 @@ from unittest.mock import patch
 
 from aiohomekit.model.characteristics import CharacteristicsTypes
 from aiohomekit.model.services import ServicesTypes
-from aiohomekit.testing import FakeController
 
 from homeassistant.components.homekit_controller.const import ENTITY_MAP
 from homeassistant.const import EVENT_HOMEASSISTANT_STOP
@@ -35,19 +34,16 @@ async def test_unload_on_stop(hass, utcnow):
 async def test_async_remove_entry(hass: HomeAssistant):
     """Test unpairing a component."""
     helper = await setup_test_component(hass, create_motion_sensor_service)
+    controller = helper.pairing.controller
 
     hkid = "00:00:00:00:00:00"
 
-    with patch("aiohomekit.Controller") as controller_cls:
-        # Setup a fake controller with 1 pairing
-        controller = controller_cls.return_value = FakeController()
-        await controller.add_paired_device([helper.accessory], hkid)
-        assert len(controller.pairings) == 1
+    assert len(controller.pairings) == 1
 
-        assert hkid in hass.data[ENTITY_MAP].storage_data
+    assert hkid in hass.data[ENTITY_MAP].storage_data
 
-        # Remove it via config entry and number of pairings should go down
-        await helper.config_entry.async_remove(hass)
-        assert len(controller.pairings) == 0
+    # Remove it via config entry and number of pairings should go down
+    await helper.config_entry.async_remove(hass)
+    assert len(controller.pairings) == 0
 
-        assert hkid not in hass.data[ENTITY_MAP].storage_data
+    assert hkid not in hass.data[ENTITY_MAP].storage_data