Fix Hue service being removed on entry reload (#48663)

homeassistant/components/hue/__init__.py

@@ -3,19 +3,18 @@ import asyncio
 import logging
 
 from aiohue.util import normalize_bridge_id
+import voluptuous as vol
 
 from homeassistant import config_entries, core
 from homeassistant.components import persistent_notification
-from homeassistant.helpers import device_registry as dr
+from homeassistant.helpers import config_validation as cv, device_registry as dr
+from homeassistant.helpers.service import verify_domain_control
 
-from .bridge import (
+from .bridge import HueBridge
+from .const import (
     ATTR_GROUP_NAME,
     ATTR_SCENE_NAME,
-    SCENE_SCHEMA,
+    ATTR_TRANSITION,
-    SERVICE_HUE_SCENE,
-    HueBridge,
-)
-from .const import (
     CONF_ALLOW_HUE_GROUPS,
     CONF_ALLOW_UNREACHABLE,
     DEFAULT_ALLOW_HUE_GROUPS,
@@ -24,46 +23,7 @@ from .const import (
 )
 
 _LOGGER = logging.getLogger(__name__)
-
+SERVICE_HUE_SCENE = "hue_activate_scene"
-
-async def async_setup(hass, config):
-    """Set up the Hue platform."""
-
-    async def hue_activate_scene(call, skip_reload=True):
-        """Handle activation of Hue scene."""
-        # Get parameters
-        group_name = call.data[ATTR_GROUP_NAME]
-        scene_name = call.data[ATTR_SCENE_NAME]
-
-        # Call the set scene function on each bridge
-        tasks = [
-            bridge.hue_activate_scene(
-                call, updated=skip_reload, hide_warnings=skip_reload
-            )
-            for bridge in hass.data[DOMAIN].values()
-            if isinstance(bridge, HueBridge)
-        ]
-        results = await asyncio.gather(*tasks)
-
-        # Did *any* bridge succeed? If not, refresh / retry
-        # Note that we'll get a "None" value for a successful call
-        if None not in results:
-            if skip_reload:
-                await hue_activate_scene(call, skip_reload=False)
-                return
-            _LOGGER.warning(
-                "No bridge was able to activate " "scene %s in group %s",
-                scene_name,
-                group_name,
-            )
-
-    # Register a local handler for scene activation
-    hass.services.async_register(
-        DOMAIN, SERVICE_HUE_SCENE, hue_activate_scene, schema=SCENE_SCHEMA
-    )
-
-    hass.data[DOMAIN] = {}
-    return True
 
 
 async def async_setup_entry(
@@ -104,7 +64,9 @@ async def async_setup_entry(
     if not await bridge.async_setup():
         return False
 
-    hass.data[DOMAIN][entry.entry_id] = bridge
+    _register_services(hass)
+
+    hass.data.setdefault(DOMAIN, {})[entry.entry_id] = bridge
     config = bridge.api.config
 
     # For backwards compat
@@ -172,5 +134,55 @@ async def async_setup_entry(
 async def async_unload_entry(hass, entry):
     """Unload a config entry."""
     bridge = hass.data[DOMAIN].pop(entry.entry_id)
+    if len(hass.data[DOMAIN]) == 0:
+        hass.data.pop(DOMAIN)
         hass.services.async_remove(DOMAIN, SERVICE_HUE_SCENE)
     return await bridge.async_reset()
+
+
+@core.callback
+def _register_services(hass):
+    """Register Hue services."""
+
+    async def hue_activate_scene(call, skip_reload=True):
+        """Handle activation of Hue scene."""
+        # Get parameters
+        group_name = call.data[ATTR_GROUP_NAME]
+        scene_name = call.data[ATTR_SCENE_NAME]
+
+        # Call the set scene function on each bridge
+        tasks = [
+            bridge.hue_activate_scene(
+                call.data, updated=skip_reload, hide_warnings=skip_reload
+            )
+            for bridge in hass.data[DOMAIN].values()
+            if isinstance(bridge, HueBridge)
+        ]
+        results = await asyncio.gather(*tasks)
+
+        # Did *any* bridge succeed? If not, refresh / retry
+        # Note that we'll get a "None" value for a successful call
+        if None not in results:
+            if skip_reload:
+                await hue_activate_scene(call, skip_reload=False)
+                return
+            _LOGGER.warning(
+                "No bridge was able to activate " "scene %s in group %s",
+                scene_name,
+                group_name,
+            )
+
+    if DOMAIN not in hass.data:
+        # Register a local handler for scene activation
+        hass.services.async_register(
+            DOMAIN,
+            SERVICE_HUE_SCENE,
+            verify_domain_control(hass, DOMAIN)(hue_activate_scene),
+            schema=vol.Schema(
+                {
+                    vol.Required(ATTR_GROUP_NAME): cv.string,
+                    vol.Required(ATTR_SCENE_NAME): cv.string,
+                    vol.Optional(ATTR_TRANSITION): cv.positive_int,
+                }
+            ),
+        )

homeassistant/components/hue/bridge.py

@@ -7,14 +7,16 @@ from aiohttp import client_exceptions
 import aiohue
 import async_timeout
 import slugify as unicode_slug
-import voluptuous as vol
 
 from homeassistant import core
 from homeassistant.const import HTTP_INTERNAL_SERVER_ERROR
 from homeassistant.exceptions import ConfigEntryNotReady
-from homeassistant.helpers import aiohttp_client, config_validation as cv
+from homeassistant.helpers import aiohttp_client
 
 from .const import (
+    ATTR_GROUP_NAME,
+    ATTR_SCENE_NAME,
+    ATTR_TRANSITION,
     CONF_ALLOW_HUE_GROUPS,
     CONF_ALLOW_UNREACHABLE,
     DEFAULT_ALLOW_HUE_GROUPS,
@@ -25,17 +27,6 @@ from .errors import AuthenticationRequired, CannotConnect
 from .helpers import create_config_flow
 from .sensor_base import SensorManager
 
-SERVICE_HUE_SCENE = "hue_activate_scene"
-ATTR_GROUP_NAME = "group_name"
-ATTR_SCENE_NAME = "scene_name"
-ATTR_TRANSITION = "transition"
-SCENE_SCHEMA = vol.Schema(
-    {
-        vol.Required(ATTR_GROUP_NAME): cv.string,
-        vol.Required(ATTR_SCENE_NAME): cv.string,
-        vol.Optional(ATTR_TRANSITION): cv.positive_int,
-    }
-)
 # How long should we sleep if the hub is busy
 HUB_BUSY_SLEEP = 0.5
 _LOGGER = logging.getLogger(__name__)
@@ -202,11 +193,11 @@ class HueBridge:
         # None and True are OK
         return False not in results
 
-    async def hue_activate_scene(self, call, updated=False, hide_warnings=False):
+    async def hue_activate_scene(self, data, skip_reload=False, hide_warnings=False):
         """Service to call directly into bridge to set scenes."""
-        group_name = call.data[ATTR_GROUP_NAME]
+        group_name = data[ATTR_GROUP_NAME]
-        scene_name = call.data[ATTR_SCENE_NAME]
+        scene_name = data[ATTR_SCENE_NAME]
-        transition = call.data.get(ATTR_TRANSITION)
+        transition = data.get(ATTR_TRANSITION)
 
         group = next(
             (group for group in self.api.groups.values() if group.name == group_name),
@@ -226,10 +217,10 @@ class HueBridge:
         )
 
         # If we can't find it, fetch latest info.
-        if not updated and (group is None or scene is None):
+        if not skip_reload and (group is None or scene is None):
             await self.async_request_call(self.api.groups.update)
             await self.async_request_call(self.api.scenes.update)
-            return await self.hue_activate_scene(call, updated=True)
+            return await self.hue_activate_scene(data, skip_reload=True)
 
         if group is None:
             if not hide_warnings:

homeassistant/components/hue/const.py

@@ -18,3 +18,7 @@ GROUP_TYPE_LIGHT_GROUP = "LightGroup"
 GROUP_TYPE_ROOM = "Room"
 GROUP_TYPE_LUMINAIRE = "Luminaire"
 GROUP_TYPE_LIGHT_SOURCE = "LightSource"
+
+ATTR_GROUP_NAME = "group_name"
+ATTR_SCENE_NAME = "scene_name"
+ATTR_TRANSITION = "transition"

tests/components/hue/test_bridge.py

@@ -185,7 +185,7 @@ async def test_hue_activate_scene(hass, mock_api):
     call = Mock()
     call.data = {"group_name": "Group 1", "scene_name": "Cozy dinner"}
     with patch("aiohue.Bridge", return_value=mock_api):
-        assert await hue_bridge.hue_activate_scene(call) is None
+        assert await hue_bridge.hue_activate_scene(call.data) is None
 
     assert len(mock_api.mock_requests) == 3
     assert mock_api.mock_requests[2]["json"]["scene"] == "scene_1"
@@ -220,7 +220,7 @@ async def test_hue_activate_scene_transition(hass, mock_api):
     call = Mock()
     call.data = {"group_name": "Group 1", "scene_name": "Cozy dinner", "transition": 30}
     with patch("aiohue.Bridge", return_value=mock_api):
-        assert await hue_bridge.hue_activate_scene(call) is None
+        assert await hue_bridge.hue_activate_scene(call.data) is None
 
     assert len(mock_api.mock_requests) == 3
     assert mock_api.mock_requests[2]["json"]["scene"] == "scene_1"
@@ -255,7 +255,7 @@ async def test_hue_activate_scene_group_not_found(hass, mock_api):
     call = Mock()
     call.data = {"group_name": "Group 1", "scene_name": "Cozy dinner"}
     with patch("aiohue.Bridge", return_value=mock_api):
-        assert await hue_bridge.hue_activate_scene(call) is False
+        assert await hue_bridge.hue_activate_scene(call.data) is False
 
 
 async def test_hue_activate_scene_scene_not_found(hass, mock_api):
@@ -285,4 +285,4 @@ async def test_hue_activate_scene_scene_not_found(hass, mock_api):
     call = Mock()
     call.data = {"group_name": "Group 1", "scene_name": "Cozy dinner"}
     with patch("aiohue.Bridge", return_value=mock_api):
-        assert await hue_bridge.hue_activate_scene(call) is False
+        assert await hue_bridge.hue_activate_scene(call.data) is False

tests/components/hue/test_init.py

@@ -27,7 +27,7 @@ async def test_setup_with_no_config(hass):
     assert len(hass.config_entries.flow.async_progress()) == 0
 
     # No configs stored
-    assert hass.data[hue.DOMAIN] == {}
+    assert hue.DOMAIN not in hass.data
 
 
 async def test_unload_entry(hass, mock_bridge_setup):
@@ -41,7 +41,7 @@ async def test_unload_entry(hass, mock_bridge_setup):
     mock_bridge_setup.async_reset = AsyncMock(return_value=True)
     assert await hue.async_unload_entry(hass, entry)
     assert len(mock_bridge_setup.async_reset.mock_calls) == 1
-    assert hass.data[hue.DOMAIN] == {}
+    assert hue.DOMAIN not in hass.data
 
 
 async def test_setting_unique_id(hass, mock_bridge_setup):

tests/components/hue/test_init_multiple_bridges.py

@@ -1,6 +1,5 @@
 """Test Hue init with multiple bridges."""
-
+from unittest.mock import AsyncMock, Mock, patch
-from unittest.mock import Mock, patch
 
 from aiohue.groups import Groups
 from aiohue.lights import Lights
@@ -13,6 +12,8 @@ from homeassistant.components import hue
 from homeassistant.components.hue import sensor_base as hue_sensor_base
 from homeassistant.setup import async_setup_component
 
+from tests.common import MockConfigEntry
+
 
 async def setup_component(hass):
     """Hue component."""
@@ -109,10 +110,9 @@ async def test_hue_activate_scene_zero_responds(
 async def setup_bridge(hass, mock_bridge, config_entry):
     """Load the Hue light platform with the provided bridge."""
     mock_bridge.config_entry = config_entry
-    hass.data[hue.DOMAIN][config_entry.entry_id] = mock_bridge
+    config_entry.add_to_hass(hass)
-    await hass.config_entries.async_forward_entry_setup(config_entry, "light")
+    with patch("homeassistant.components.hue.HueBridge", return_value=mock_bridge):
-    # To flush out the service call to update the group
+        await hass.config_entries.async_setup(config_entry.entry_id)
-    await hass.async_block_till_done()
 
 
 @pytest.fixture
@@ -129,14 +129,10 @@ def mock_config_entry2(hass):
 
 def create_config_entry():
     """Mock a config entry."""
-    return config_entries.ConfigEntry(
+    return MockConfigEntry(
-        1,
+        domain=hue.DOMAIN,
-        hue.DOMAIN,
+        data={"host": "mock-host"},
-        "Mock Title",
+        connection_class=config_entries.CONN_CLASS_LOCAL_POLL,
-        {"host": "mock-host"},
-        "test",
-        config_entries.CONN_CLASS_LOCAL_POLL,
-        system_options={},
     )
 
 
@@ -163,6 +159,7 @@ def create_mock_bridge(hass):
         api=Mock(),
         reset_jobs=[],
         spec=hue.HueBridge,
+        async_setup=AsyncMock(return_value=True),
     )
     bridge.sensor_manager = hue_sensor_base.SensorManager(bridge)
     bridge.mock_requests = []