From 2976bbbbdd91ec4276ababf0dd8ed367e8e4b3ba Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Thu, 20 May 2021 00:08:23 -0700
Subject: [PATCH] Store Hue bridge in hass.data before setting up platforms
 (#50703)

* Store bridge in hass.data before setting up platforms

* Self
---
 homeassistant/components/hue/__init__.py |  7 +++----
 homeassistant/components/hue/bridge.py   |  9 +++++++-
 tests/components/hue/conftest.py         | 26 +++++++++++++++++++++++-
 tests/components/hue/test_init.py        |  9 ++++++--
 4 files changed, 43 insertions(+), 8 deletions(-)

diff --git a/homeassistant/components/hue/__init__.py b/homeassistant/components/hue/__init__.py
index 6bbe3d9ebdd..71b62e22d33 100644
--- a/homeassistant/components/hue/__init__.py
+++ b/homeassistant/components/hue/__init__.py
@@ -66,7 +66,6 @@ async def async_setup_entry(
 
     _register_services(hass)
 
-    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = bridge
     config = bridge.api.config
 
     # For backwards compat
@@ -133,11 +132,11 @@ async def async_setup_entry(
 
 async def async_unload_entry(hass, entry):
     """Unload a config entry."""
-    bridge = hass.data[DOMAIN].pop(entry.entry_id)
+    unload_success = await hass.data[DOMAIN][entry.entry_id].async_reset()
     if len(hass.data[DOMAIN]) == 0:
         hass.data.pop(DOMAIN)
         hass.services.async_remove(DOMAIN, SERVICE_HUE_SCENE)
-    return await bridge.async_reset()
+    return unload_success
 
 
 @core.callback
@@ -172,7 +171,7 @@ def _register_services(hass):
                 group_name,
             )
 
-    if DOMAIN not in hass.data:
+    if not hass.services.has_service(DOMAIN, SERVICE_HUE_SCENE):
         # Register a local handler for scene activation
         hass.services.async_register(
             DOMAIN,
diff --git a/homeassistant/components/hue/bridge.py b/homeassistant/components/hue/bridge.py
index 8f027aee033..1f19138b28f 100644
--- a/homeassistant/components/hue/bridge.py
+++ b/homeassistant/components/hue/bridge.py
@@ -23,6 +23,7 @@ from .const import (
     CONF_ALLOW_UNREACHABLE,
     DEFAULT_ALLOW_HUE_GROUPS,
     DEFAULT_ALLOW_UNREACHABLE,
+    DOMAIN,
     LOGGER,
 )
 from .errors import AuthenticationRequired, CannotConnect
@@ -107,6 +108,7 @@ class HueBridge:
         if bridge.sensors is not None:
             self.sensor_manager = SensorManager(self)
 
+        hass.data.setdefault(DOMAIN, {})[self.config_entry.entry_id] = self
         hass.config_entries.async_setup_platforms(self.config_entry, PLATFORMS)
 
         self.parallel_updates_semaphore = asyncio.Semaphore(
@@ -173,10 +175,15 @@ class HueBridge:
 
         # If setup was successful, we set api variable, forwarded entry and
         # register service
-        return await self.hass.config_entries.async_unload_platforms(
+        unload_success = await self.hass.config_entries.async_unload_platforms(
             self.config_entry, PLATFORMS
         )
 
+        if unload_success:
+            self.hass.data[DOMAIN].pop(self.config_entry.entry_id)
+
+        return unload_success
+
     async def hue_activate_scene(self, data, skip_reload=False, hide_warnings=False):
         """Service to call directly into bridge to set scenes."""
         if self.api.scenes is None:
diff --git a/tests/components/hue/conftest.py b/tests/components/hue/conftest.py
index 3aecacac58d..b5c2aec3042 100644
--- a/tests/components/hue/conftest.py
+++ b/tests/components/hue/conftest.py
@@ -32,6 +32,7 @@ def create_mock_bridge(hass):
         allow_unreachable=False,
         allow_groups=False,
         api=create_mock_api(hass),
+        config_entry=None,
         reset_jobs=[],
         spec=hue.HueBridge,
     )
@@ -41,10 +42,25 @@ def create_mock_bridge(hass):
     bridge.mock_group_responses = bridge.api.mock_group_responses
     bridge.mock_sensor_responses = bridge.api.mock_sensor_responses
 
+    async def async_setup():
+        if bridge.config_entry:
+            hass.data.setdefault(hue.DOMAIN, {})[bridge.config_entry.entry_id] = bridge
+        return True
+
+    bridge.async_setup = async_setup
+
     async def async_request_call(task):
         await task()
 
     bridge.async_request_call = async_request_call
+
+    async def async_reset():
+        if bridge.config_entry:
+            hass.data[hue.DOMAIN].pop(bridge.config_entry.entry_id)
+        return True
+
+    bridge.async_reset = async_reset
+
     return bridge
 
 
@@ -80,7 +96,15 @@ def create_mock_api(hass):
 
     logger = logging.getLogger(__name__)
 
-    api.config.apiversion = "9.9.9"
+    api.config = Mock(
+        bridgeid="ff:ff:ff:ff:ff:ff",
+        mac="aa:bb:cc:dd:ee:ff",
+        modelid="BSB002",
+        apiversion="9.9.9",
+        swversion="1935144040",
+    )
+    api.config.name = "Home"
+
     api.lights = Lights(logger, {}, [], mock_request)
     api.groups = Groups(logger, {}, [], mock_request)
     api.sensors = Sensors(logger, {}, [], mock_request)
diff --git a/tests/components/hue/test_init.py b/tests/components/hue/test_init.py
index 0c1d75c2ce2..afc920a6667 100644
--- a/tests/components/hue/test_init.py
+++ b/tests/components/hue/test_init.py
@@ -38,9 +38,14 @@ async def test_unload_entry(hass, mock_bridge_setup):
     assert await async_setup_component(hass, hue.DOMAIN, {}) is True
     assert len(mock_bridge_setup.mock_calls) == 1
 
-    mock_bridge_setup.async_reset = AsyncMock(return_value=True)
+    hass.data[hue.DOMAIN] = {entry.entry_id: mock_bridge_setup}
+
+    async def mock_reset():
+        hass.data[hue.DOMAIN].pop(entry.entry_id)
+        return True
+
+    mock_bridge_setup.async_reset = mock_reset
     assert await hue.async_unload_entry(hass, entry)
-    assert len(mock_bridge_setup.async_reset.mock_calls) == 1
     assert hue.DOMAIN not in hass.data