Merge pull request #229 from azogue/feature/refactor

Fix setup flow broken with last refactor

custom_components/huesensor/data_manager.py

@@ -2,7 +2,7 @@
 import asyncio
 from datetime import timedelta
 import logging
-from typing import AsyncIterable, List, Tuple
+from typing import AsyncIterable, Set, Tuple
 
 from homeassistant.components.hue import DOMAIN as HUE_DOMAIN, HueBridge
 from homeassistant.helpers.entity import Entity
@@ -25,13 +25,11 @@
 DEFAULT_SCAN_INTERVAL = timedelta(seconds=0.5)
 
 
-def get_bridges(hass) -> List[HueBridge]:
+async def async_get_bridges(hass) -> AsyncIterable[HueBridge]:
     """Retrieve Hue bridges from loaded official Hue integration."""
-    return [
-        entry
-        for entry in hass.data[HUE_DOMAIN].values()
-        if isinstance(entry, HueBridge) and entry.api
-    ]
+    for entry in hass.data[HUE_DOMAIN].values():
+        if isinstance(entry, HueBridge) and entry.api:
+            yield entry
 
 
 class HueSensorData:
@@ -43,23 +41,24 @@ def __init__(self, hass):
         self.lock = asyncio.Lock()
         self.data = {}
         self.sensors = {}
+        self.registered_entities = {}
         self.available = False
         self._scan_interval = None
         self._update_listener = None
 
-    async def _iter_data(self) -> AsyncIterable[Tuple[bool, str, str, dict]]:
-        bridges = get_bridges(self.hass)
-        await asyncio.gather(
-            *[
-                bridge.sensor_manager.coordinator.async_request_refresh()
-                for bridge in bridges
-            ]
-        )
-        for bridge in bridges:
+        # delayed setup and discovery with platform + model filter
+        self._registered_models: Set[str] = set()
+        self._registered_platforms = {}
+
+    async def _iter_data(
+        self, models_filter: Tuple[str] = _KNOWN_MODEL_IDS
+    ) -> AsyncIterable[Tuple[bool, str, str, dict]]:
+        async for bridge in async_get_bridges(self.hass):
+            await bridge.sensor_manager.coordinator.async_request_refresh()
             data = parse_hue_api_response(
                 sensor.raw
                 for sensor in bridge.api.sensors.values()
-                if sensor.raw["modelid"].startswith(_KNOWN_MODEL_IDS)
+                if sensor.raw["modelid"].startswith(models_filter)
             )
             for dev_id, dev_data in data.items():
                 updated = False
@@ -111,19 +110,25 @@ async def async_stop_scheduler(self):
             self.available = False
         _LOGGER.debug(f"Stopped polling with {self._scan_interval}")
 
-    async def async_add_platform_entities(
-        self, entity_cls, platform_models, async_add_entities, scan_interval,
-    ):
-        """Add sensor entities from platform setups."""
-        new_entities = []
-        async for updated, model, dev_id, dev_data in self._iter_data():
-            if model in platform_models and dev_id not in self.sensors:
-                platform_entity = entity_cls(dev_id, self)
-                self.sensors[dev_id] = platform_entity
-                new_entities.append(platform_entity)
+    def _register_new_entity(self, dev_id, model, new_entities_to_add):
+        """Register a new Entity and add it in platform queue for HA setup."""
+        # Create platform entity and register the device
+        entity_cls, async_add_entities = self._registered_platforms[model]
+        platform_entity = entity_cls(dev_id, self)
+        self.registered_entities[dev_id] = platform_entity
+
+        # Add entity to platform queue to add it to HA
+        if async_add_entities not in new_entities_to_add:
+            new_entities_to_add[async_add_entities] = [entity_cls, []]
+        new_entities_to_add[async_add_entities][1].append(platform_entity)
 
-        if new_entities:
-            async_add_entities(new_entities, True)
+    async def _add_new_entities(self, new_entities, scan_interval=None):
+        """Call HA add_entities for each platform with its discovered items."""
+        for func_add_entities, (entity_cls, entities) in new_entities.items():
+            func_add_entities(entities, True)
+
+            if scan_interval is None:
+                continue
 
             if self._scan_interval is None:
                 self._scan_interval = scan_interval
@@ -143,17 +148,47 @@ async def async_add_platform_entities(
                 await self.async_stop_scheduler()
                 self._scan_interval = scan_interval
 
+    async def async_add_platform_entities(
+        self, entity_cls, platform_models, func_add_entities, scan_interval,
+    ):
+        """Add sensor entities from platform setups."""
+        for model in platform_models:
+            self._registered_platforms[model] = (entity_cls, func_add_entities)
+            self._registered_models.add(model)
+
+        new_entities_to_add = {}
+        async for is_new, model, dev_id, _ in self._iter_data(platform_models):
+            if is_new and dev_id not in self.registered_entities:
+                self._register_new_entity(dev_id, model, new_entities_to_add)
+
+        await self._add_new_entities(new_entities_to_add, scan_interval)
+
     async def async_update_from_bridges(self, now=None):
         """Request data from bridges and update sensors data."""
-        async for updated, _, dev_id, _ in self._iter_data():
-            if updated:
+        new_entities_to_add = {}
+        async for updated, model, dev_id, _dev_data in self._iter_data(
+            tuple(self._registered_models)
+        ):
+            if updated and dev_id not in self.registered_entities:
+                # Discovery of newly added devices
+                _LOGGER.warning(
+                    "New device discovered %s:%s. Adding it now", model, dev_id
+                )
+                self._register_new_entity(dev_id, model, new_entities_to_add)
+            elif updated and dev_id not in self.sensors:
+                # device is registered, but it is not added to hass yet ¿?
+                _LOGGER.warning(
+                    "Device %s:%s registered but not added yet", model, dev_id
+                )
+            elif updated:
                 self.sensors[dev_id].async_write_ha_state()
                 _LOGGER.debug(
                     "%s (%s): updated with state=%s",
                     self.sensors[dev_id].entity_id,
                     dev_id,
                     self.sensors[dev_id].state,
                 )
+        await self._add_new_entities(new_entities_to_add)
 
 
 class HueSensorBaseDevice(Entity):
@@ -165,23 +200,18 @@ def __init__(self, hue_id, data_manager: HueSensorData):
         self._data_manager = data_manager
 
     async def async_added_to_hass(self):
-        """When entity is added to hass."""
-        # TODO use bridge.sensor_manager.coordinator.async_add_listener
-        # self.bridge.sensor_manager.coordinator.async_add_listener(
-        #     self.async_write_ha_state
-        # )
+        """Register sensor when entity is added to hass and start updating."""
+        self._data_manager.sensors[self.unique_id] = self
         await self._data_manager.async_start_scheduler()
         _LOGGER.debug(
-            "Setup complete for %s:%s", self.__class__.__name__, self._hue_id
+            "Setup complete for %s:%s", self.__class__.__name__, self.unique_id
         )
 
     async def async_will_remove_from_hass(self):
         """When entity will be removed from hass."""
-        # self.bridge.sensor_manager.coordinator.async_remove_listener(
-        #     self.async_write_ha_state
-        # )
         _LOGGER.debug("%s: Removing entity from HA", self.entity_id)
-        self._data_manager.sensors.pop(self._hue_id)
+        self._data_manager.sensors.pop(self.unique_id)
+        self._data_manager.registered_entities.pop(self.unique_id)
 
         if self._data_manager.sensors:
             return
@@ -191,7 +221,7 @@ async def async_will_remove_from_hass(self):
     @property
     def sensor_data(self) -> dict:
         """Access to parsed sensor data."""
-        return self._data_manager.data.get(self._hue_id)
+        return self._data_manager.data.get(self.unique_id)
 
     @property
     def should_poll(self):

custom_components/huesensor/device_tracker.py

@@ -20,7 +20,7 @@
 from homeassistant.helpers.event import async_track_time_interval
 from homeassistant.util import slugify
 
-from .data_manager import get_bridges
+from .data_manager import async_get_bridges
 
 _LOGGER = logging.getLogger(__name__)
 
@@ -87,19 +87,12 @@ async def async_see_sensor(self, sensor):
 
     async def async_update_info(self, now=None):
         """Get the bridge info."""
-        bridges = get_bridges(self.hass)
-        if not bridges:
-            return
-
-        for bridge in bridges:
+        async for bridge in async_get_bridges(self.hass):
             await bridge.sensor_manager.coordinator.async_request_refresh()
-
-        sensors = [
-            self.async_see_sensor(sensor)
-            for bridge in bridges
-            for sensor in bridge.api.sensors.values()
-            if sensor.type == TYPE_GEOFENCE
-        ]
-        if not sensors:
-            return
-        await asyncio.wait(sensors)
+            tasks = [
+                self.async_see_sensor(sensor)
+                for sensor in bridge.api.sensors.values()
+                if sensor.type == TYPE_GEOFENCE
+            ]
+            if tasks:
+                await asyncio.wait(tasks)

custom_components/huesensor/hue_api_response.py

@@ -4,8 +4,8 @@
 
 from homeassistant.const import STATE_OFF, STATE_ON
 
-REMOTE_MODELS = ["RWL", "ROM", "FOH", "ZGP", "Z3-"]
-BINARY_SENSOR_MODELS = ["SML"]
+REMOTE_MODELS = ("RWL", "ROM", "FOH", "ZGP", "Z3-")
+BINARY_SENSOR_MODELS = ("SML",)
 ENTITY_ATTRS = {
     "RWL": ["last_updated", "last_button_event", "battery", "on", "reachable"],
     "ROM": ["last_updated", "last_button_event", "battery", "on", "reachable"],

tests/conftest.py

@@ -8,8 +8,14 @@
 from homeassistant.components.hue import DOMAIN as HUE_DOMAIN, HueBridge
 from homeassistant.components.hue.sensor_base import SensorManager
 from homeassistant.helpers.update_coordinator import DataUpdateCoordinator
+from homeassistant.util import slugify
 import pytest
 
+from custom_components.huesensor.data_manager import (
+    BINARY_SENSOR_MODELS,
+    HueSensorBaseDevice,
+    HueSensorData,
+)
 from .sensor_samples import (
     MOCK_GEOFENCE,
     MOCK_ZGP,
@@ -19,9 +25,24 @@
 )
 
 DEV_ID_REMOTE_1 = "ZGP_00:44:23:08"
+DEV_ID_REMOTE_2 = "RWL_00:17:88:01:10:3e:3a:dc-02"
 DEV_ID_SENSOR_1 = "SML_00:17:88:01:02:00:af:28-02"
 
 
+async def entity_test_added_to_hass(
+    data_manager: HueSensorData, entity: HueSensorBaseDevice,
+):
+    """Test routine to mock the internals of async_added_to_hass."""
+    entity.hass = data_manager.hass
+    if entity.unique_id.startswith(BINARY_SENSOR_MODELS):
+        entity.entity_id = f"binary_sensor.test_{slugify(entity.name)}"
+    else:
+        entity.entity_id = f"remote.test_{slugify(entity.name)}"
+    await entity.async_added_to_hass()
+    assert data_manager.available
+    assert entity.unique_id in data_manager.sensors
+
+
 class MockAsyncCounter:
     """
     Call counter for the hue data coordinator.
@@ -47,14 +68,20 @@ def call_count(self) -> int:
         return self._counter
 
 
+def add_sensor_data_to_bridge(bridge, sensor_key, raw_data):
+    """Append a sensor raw data packed to the mocked bridge."""
+    bridge.sensors[sensor_key] = GenericSensor(
+        raw_data["uniqueid"], deepcopy(raw_data), None
+    )
+
+
 def _make_mock_bridge(idx_bridge, *sensors):
     bridge = MagicMock(spec=Bridge)
-    bridge.sensors = {
-        f"{raw_data['type']}_{idx_bridge}_{i}": GenericSensor(
-            raw_data["uniqueid"], deepcopy(raw_data), None
+    bridge.sensors = {}
+    for i, raw_data in enumerate(sensors):
+        add_sensor_data_to_bridge(
+            bridge, f"{raw_data['type']}_{idx_bridge}_{i}", raw_data
         )
-        for i, raw_data in enumerate(sensors)
-    }
     return bridge
 
 

tests/test_binary_sensor.py

@@ -18,6 +18,7 @@
 from .conftest import (
     DEV_ID_REMOTE_1,
     DEV_ID_SENSOR_1,
+    entity_test_added_to_hass,
     patch_async_track_time_interval,
 )
 from .sensor_samples import (
@@ -80,16 +81,15 @@ async def test_platform_binary_sensor_setup(mock_hass, caplog):
             assert DOMAIN in mock_hass.data
             data_manager = mock_hass.data[DOMAIN]
             assert isinstance(data_manager, HueSensorData)
+            assert len(data_manager.registered_entities) == 1
             assert data_manager._scan_interval == timedelta(seconds=2)
-            assert len(data_manager.data) == 2
-            assert DEV_ID_REMOTE_1 in data_manager.data
+            assert len(data_manager.data) == 1
+            assert DEV_ID_REMOTE_1 not in data_manager.data
             assert DEV_ID_SENSOR_1 in data_manager.data
+            assert len(data_manager.sensors) == 0
 
-            assert len(data_manager.data) == 2
-            assert len(data_manager.sensors) == 1
-
-            assert DEV_ID_SENSOR_1 in data_manager.sensors
-            bin_sensor = data_manager.sensors[DEV_ID_SENSOR_1]
+            assert DEV_ID_SENSOR_1 in data_manager.registered_entities
+            bin_sensor = data_manager.registered_entities[DEV_ID_SENSOR_1]
             assert isinstance(bin_sensor, HueBinarySensor)
             assert bin_sensor.device_class == "motion"
             assert not bin_sensor.is_on
@@ -98,5 +98,8 @@ async def test_platform_binary_sensor_setup(mock_hass, caplog):
             assert "light_level" in bin_sensor.device_state_attributes
             assert bin_sensor.unique_id == DEV_ID_SENSOR_1
 
-            await bin_sensor.async_added_to_hass()
+            await entity_test_added_to_hass(data_manager, bin_sensor)
+            assert len(data_manager.sensors) == 1
             await bin_sensor.async_will_remove_from_hass()
+            assert len(data_manager.sensors) == 0
+            assert len(data_manager.registered_entities) == 0