From e09234ffae2ffc3f0d0f98d440ec11f72e4fe4bf Mon Sep 17 00:00:00 2001 From: zewelor Date: Thu, 10 Dec 2020 09:54:10 +0100 Subject: [PATCH] Fix yeelight unavailbility (#44061) --- homeassistant/components/yeelight/__init__.py | 66 ++++++++++++------- tests/components/yeelight/__init__.py | 3 +- tests/components/yeelight/test_init.py | 46 ++++++++++++- 3 files changed, 88 insertions(+), 27 deletions(-) diff --git a/homeassistant/components/yeelight/__init__.py b/homeassistant/components/yeelight/__init__.py index ae9d75de54f..324999c7124 100644 --- a/homeassistant/components/yeelight/__init__.py +++ b/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) diff --git a/tests/components/yeelight/__init__.py b/tests/components/yeelight/__init__.py index 9f811586a77..5405b69490b 100644 --- a/tests/components/yeelight/__init__.py +++ b/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" diff --git a/tests/components/yeelight/test_init.py b/tests/components/yeelight/test_init.py index d9c23cfa1a7..882f9944ca1 100644 --- a/tests/components/yeelight/test_init.py +++ b/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