Fix yeelight unavailbility (#44061)

homeassistant/components/yeelight/__init__.py

@@ -17,7 +17,7 @@ from homeassistant.const import (
 )
 from homeassistant.core import HomeAssistant, callback
 import homeassistant.helpers.config_validation as cv
-from homeassistant.helpers.dispatcher import dispatcher_send
+from homeassistant.helpers.dispatcher import async_dispatcher_connect, dispatcher_send
 from homeassistant.helpers.entity import Entity
 from homeassistant.helpers.event import async_track_time_interval
 
@@ -26,7 +26,7 @@ _LOGGER = logging.getLogger(__name__)
 DOMAIN = "yeelight"
 DATA_YEELIGHT = DOMAIN
 DATA_UPDATED = "yeelight_{}_data_updated"
-DEVICE_INITIALIZED = f"{DOMAIN}_device_initialized"
+DEVICE_INITIALIZED = "yeelight_{}_device_initialized"
 
 DEFAULT_NAME = "Yeelight"
 DEFAULT_TRANSITION = 350
@@ -181,8 +181,19 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     """Set up Yeelight from a config entry."""
 
     async def _initialize(host: str, capabilities: Optional[dict] = None) -> None:
-        device = await _async_setup_device(hass, host, entry, capabilities)
+        async_dispatcher_connect(
+            hass,
+            DEVICE_INITIALIZED.format(host),
+            _load_platforms,
+        )
+
+        device = await _async_get_device(hass, host, entry, capabilities)
         hass.data[DOMAIN][DATA_CONFIG_ENTRIES][entry.entry_id][DATA_DEVICE] = device
+
+        await device.async_setup()
+
+    async def _load_platforms():
+
         for component in PLATFORMS:
             hass.async_create_task(
                 hass.config_entries.async_forward_entry_setup(entry, component)
@@ -249,28 +260,6 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry):
     return unload_ok
 
 
-async def _async_setup_device(
-    hass: HomeAssistant,
-    host: str,
-    entry: ConfigEntry,
-    capabilities: Optional[dict],
-) -> None:
-    # Get model from config and capabilities
-    model = entry.options.get(CONF_MODEL)
-    if not model and capabilities is not None:
-        model = capabilities.get("model")
-
-    # Set up device
-    bulb = Bulb(host, model=model or None)
-    if capabilities is None:
-        capabilities = await hass.async_add_executor_job(bulb.get_capabilities)
-
-    device = YeelightDevice(hass, host, entry.options, bulb, capabilities)
-    await hass.async_add_executor_job(device.update)
-    await device.async_setup()
-    return device
-
-
 @callback
 def _async_unique_name(capabilities: dict) -> str:
     """Generate name from capabilities."""
@@ -374,6 +363,7 @@ class YeelightDevice:
         self._device_type = None
         self._available = False
         self._remove_time_tracker = None
+        self._initialized = False
 
         self._name = host  # Default name is host
         if capabilities:
@@ -495,6 +485,8 @@ class YeelightDevice:
         try:
             self.bulb.get_properties(UPDATE_REQUEST_PROPERTIES)
             self._available = True
+            if not self._initialized:
+                self._initialize_device()
         except BulbException as ex:
             if self._available:  # just inform once
                 _LOGGER.error(
@@ -522,6 +514,11 @@ class YeelightDevice:
                 ex,
             )
 
+    def _initialize_device(self):
+        self._get_capabilities()
+        self._initialized = True
+        dispatcher_send(self._hass, DEVICE_INITIALIZED.format(self._host))
+
     def update(self):
         """Update device properties and send data updated signal."""
         self._update_properties()
@@ -584,3 +581,22 @@ class YeelightEntity(Entity):
     def update(self) -> None:
         """Update the entity."""
         self._device.update()
+
+
+async def _async_get_device(
+    hass: HomeAssistant,
+    host: str,
+    entry: ConfigEntry,
+    capabilities: Optional[dict],
+) -> YeelightDevice:
+    # Get model from config and capabilities
+    model = entry.options.get(CONF_MODEL)
+    if not model and capabilities is not None:
+        model = capabilities.get("model")
+
+    # Set up device
+    bulb = Bulb(host, model=model or None)
+    if capabilities is None:
+        capabilities = await hass.async_add_executor_job(bulb.get_capabilities)
+
+    return YeelightDevice(hass, host, entry.options, bulb, capabilities)

tests/components/yeelight/__init__.py

@@ -55,7 +55,8 @@ PROPERTIES = {
     "current_brightness": "30",
 }
 
-ENTITY_BINARY_SENSOR = f"binary_sensor.{UNIQUE_NAME}_nightlight"
+ENTITY_BINARY_SENSOR_TEMPLATE = "binary_sensor.{}_nightlight"
+ENTITY_BINARY_SENSOR = ENTITY_BINARY_SENSOR_TEMPLATE.format(UNIQUE_NAME)
 ENTITY_LIGHT = f"light.{UNIQUE_NAME}"
 ENTITY_NIGHTLIGHT = f"light.{UNIQUE_NAME}_nightlight"
 ENTITY_AMBILIGHT = f"light.{UNIQUE_NAME}_ambilight"

tests/components/yeelight/test_init.py

@@ -1,21 +1,27 @@
 """Test Yeelight."""
+from unittest.mock import MagicMock
+
 from yeelight import BulbType
 
 from homeassistant.components.yeelight import (
     CONF_NIGHTLIGHT_SWITCH,
     CONF_NIGHTLIGHT_SWITCH_TYPE,
+    DATA_CONFIG_ENTRIES,
+    DATA_DEVICE,
     DOMAIN,
     NIGHTLIGHT_SWITCH_TYPE_LIGHT,
 )
-from homeassistant.const import CONF_DEVICES, CONF_NAME
+from homeassistant.const import CONF_DEVICES, CONF_HOST, CONF_NAME
 from homeassistant.core import HomeAssistant
 from homeassistant.helpers import entity_registry
 from homeassistant.setup import async_setup_component
 
 from . import (
+    CAPABILITIES,
     CONFIG_ENTRY_DATA,
     ENTITY_AMBILIGHT,
     ENTITY_BINARY_SENSOR,
+    ENTITY_BINARY_SENSOR_TEMPLATE,
     ENTITY_LIGHT,
     ENTITY_NIGHTLIGHT,
     ID,
@@ -115,6 +121,7 @@ async def test_unique_ids_entry(hass: HomeAssistant):
 
     mocked_bulb = _mocked_bulb()
     mocked_bulb.bulb_type = BulbType.WhiteTempMood
+
     with _patch_discovery(MODULE), patch(f"{MODULE}.Bulb", return_value=mocked_bulb):
         assert await hass.config_entries.async_setup(config_entry.entry_id)
         await hass.async_block_till_done()
@@ -132,3 +139,40 @@ async def test_unique_ids_entry(hass: HomeAssistant):
     assert (
         er.async_get(ENTITY_AMBILIGHT).unique_id == f"{config_entry.entry_id}-ambilight"
     )
+
+
+async def test_bulb_off_while_adding_in_ha(hass: HomeAssistant):
+    """Test Yeelight off while adding to ha, for example on HA start."""
+    config_entry = MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            **CONFIG_ENTRY_DATA,
+            CONF_HOST: IP_ADDRESS,
+        },
+        unique_id=ID,
+    )
+    config_entry.add_to_hass(hass)
+
+    mocked_bulb = _mocked_bulb(True)
+    mocked_bulb.bulb_type = BulbType.WhiteTempMood
+
+    with patch(f"{MODULE}.Bulb", return_value=mocked_bulb), patch(
+        f"{MODULE}.config_flow.yeelight.Bulb", return_value=mocked_bulb
+    ):
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+
+    binary_sensor_entity_id = ENTITY_BINARY_SENSOR_TEMPLATE.format(
+        IP_ADDRESS.replace(".", "_")
+    )
+    er = await entity_registry.async_get_registry(hass)
+    assert er.async_get(binary_sensor_entity_id) is None
+
+    type(mocked_bulb).get_capabilities = MagicMock(CAPABILITIES)
+    type(mocked_bulb).get_properties = MagicMock(None)
+
+    hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][DATA_DEVICE].update()
+    await hass.async_block_till_done()
+
+    er = await entity_registry.async_get_registry(hass)
+    assert er.async_get(binary_sensor_entity_id) is not None