diff --git a/homeassistant/components/yeelight/__init__.py b/homeassistant/components/yeelight/__init__.py index c9e654bca9a..0ea4eb8e84f 100644 --- a/homeassistant/components/yeelight/__init__.py +++ b/homeassistant/components/yeelight/__init__.py @@ -2,13 +2,17 @@ from __future__ import annotations import asyncio +import contextlib from datetime import timedelta import logging +from urllib.parse import urlparse +from async_upnp_client.search import SSDPListener import voluptuous as vol -from yeelight import BulbException, discover_bulbs +from yeelight import BulbException from yeelight.aio import KEY_CONNECTED, AsyncBulb +from homeassistant import config_entries from homeassistant.config_entries import SOURCE_IMPORT, ConfigEntry, ConfigEntryNotReady from homeassistant.const import ( CONF_DEVICES, @@ -24,6 +28,7 @@ from homeassistant.helpers.dispatcher import ( async_dispatcher_send, ) from homeassistant.helpers.entity import DeviceInfo, Entity +from homeassistant.helpers.event import async_call_later, async_track_time_interval from homeassistant.helpers.typing import ConfigType _LOGGER = logging.getLogger(__name__) @@ -69,6 +74,12 @@ ACTIVE_COLOR_FLOWING = "1" NIGHTLIGHT_SWITCH_TYPE_LIGHT = "light" DISCOVERY_INTERVAL = timedelta(seconds=60) +SSDP_TARGET = ("239.255.255.250", 1982) +SSDP_ST = "wifi_bulb" +DISCOVERY_ATTEMPTS = 3 +DISCOVERY_SEARCH_INTERVAL = timedelta(seconds=2) +DISCOVERY_TIMEOUT = 2 + YEELIGHT_RGB_TRANSITION = "RGBTransition" YEELIGHT_HSV_TRANSACTION = "HSVTransition" @@ -193,20 +204,12 @@ async def _async_initialize( hass.config_entries.async_setup_platforms(entry, PLATFORMS) if not device: + # get device and start listening for local pushes device = await _async_get_device(hass, host, entry) + + await device.async_setup() entry_data[DATA_DEVICE] = device - # start listening for local pushes - await device.bulb.async_listen(device.async_update_callback) - - # register stop callback to shutdown listening for local pushes - async def async_stop_listen_task(event): - """Stop listen thread.""" - _LOGGER.debug("Shutting down Yeelight Listener") - await device.bulb.async_stop_listening() - - hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_stop_listen_task) - entry.async_on_unload( async_dispatcher_connect( hass, DEVICE_INITIALIZED.format(host), _async_load_platforms @@ -251,7 +254,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: if entry.data.get(CONF_HOST): try: device = await _async_get_device(hass, entry.data[CONF_HOST], entry) - except OSError as ex: + except BulbException as ex: # If CONF_ID is not valid we cannot fallback to discovery # so we must retry by raising ConfigEntryNotReady if not entry.data.get(CONF_ID): @@ -267,16 +270,16 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: raise ConfigEntryNotReady from ex return True - # discovery - scanner = YeelightScanner.async_get(hass) - - async def _async_from_discovery(host: str) -> None: + async def _async_from_discovery(capabilities: dict[str, str]) -> None: + host = urlparse(capabilities["location"]).hostname try: await _async_initialize(hass, entry, host) except BulbException: _LOGGER.exception("Failed to connect to bulb at %s", host) - scanner.async_register_callback(entry.data[CONF_ID], _async_from_discovery) + # discovery + scanner = YeelightScanner.async_get(hass) + await scanner.async_register_callback(entry.data[CONF_ID], _async_from_discovery) return True @@ -294,10 +297,11 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: scanner = YeelightScanner.async_get(hass) scanner.async_unregister_callback(entry.data[CONF_ID]) - device = entry_data[DATA_DEVICE] - _LOGGER.debug("Shutting down Yeelight Listener") - await device.bulb.async_stop_listening() - _LOGGER.debug("Yeelight Listener stopped") + if DATA_DEVICE in entry_data: + device = entry_data[DATA_DEVICE] + _LOGGER.debug("Shutting down Yeelight Listener") + await device.bulb.async_stop_listening() + _LOGGER.debug("Yeelight Listener stopped") data_config_entries.pop(entry.entry_id) @@ -307,9 +311,9 @@ async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: @callback def _async_unique_name(capabilities: dict) -> str: """Generate name from capabilities.""" - model = capabilities["model"] - unique_id = capabilities["id"] - return f"yeelight_{model}_{unique_id}" + model = str(capabilities["model"]).replace("_", " ").title() + short_id = hex(int(capabilities["id"], 16)) + return f"Yeelight {model} {short_id}" async def _async_update_listener(hass: HomeAssistant, entry: ConfigEntry): @@ -333,88 +337,147 @@ class YeelightScanner: def __init__(self, hass: HomeAssistant) -> None: """Initialize class.""" self._hass = hass - self._seen = {} self._callbacks = {} - self._scan_task = None + self._host_discovered_events = {} + self._unique_id_capabilities = {} + self._host_capabilities = {} + self._track_interval = None + self._listener = None + self._connected_event = None - async def _async_scan(self): - _LOGGER.debug("Yeelight scanning") - # Run 3 times as packets can get lost - for _ in range(3): - devices = await self._hass.async_add_executor_job(discover_bulbs) - for device in devices: - unique_id = device["capabilities"]["id"] - if unique_id in self._seen: - continue - host = device["ip"] - self._seen[unique_id] = host - _LOGGER.debug("Yeelight discovered at %s", host) - if unique_id in self._callbacks: - self._hass.async_create_task(self._callbacks[unique_id](host)) - self._callbacks.pop(unique_id) - if len(self._callbacks) == 0: - self._async_stop_scan() + async def async_setup(self): + """Set up the scanner.""" + if self._connected_event: + await self._connected_event.wait() + return + self._connected_event = asyncio.Event() - await asyncio.sleep(DISCOVERY_INTERVAL.total_seconds()) - self._scan_task = self._hass.loop.create_task(self._async_scan()) + async def _async_connected(): + self._listener.async_search() + self._connected_event.set() + + self._listener = SSDPListener( + async_callback=self._async_process_entry, + service_type=SSDP_ST, + target=SSDP_TARGET, + async_connect_callback=_async_connected, + ) + await self._listener.async_start() + await self._connected_event.wait() + + async def async_discover(self): + """Discover bulbs.""" + await self.async_setup() + for _ in range(DISCOVERY_ATTEMPTS): + self._listener.async_search() + await asyncio.sleep(DISCOVERY_SEARCH_INTERVAL.total_seconds()) + return self._unique_id_capabilities.values() @callback - def _async_start_scan(self): + def async_scan(self, *_): + """Send discovery packets.""" + _LOGGER.debug("Yeelight scanning") + self._listener.async_search() + + async def async_get_capabilities(self, host): + """Get capabilities via SSDP.""" + if host in self._host_capabilities: + return self._host_capabilities[host] + + host_event = asyncio.Event() + self._host_discovered_events.setdefault(host, []).append(host_event) + await self.async_setup() + + self._listener.async_search((host, SSDP_TARGET[1])) + + with contextlib.suppress(asyncio.TimeoutError): + await asyncio.wait_for(host_event.wait(), timeout=DISCOVERY_TIMEOUT) + + self._host_discovered_events[host].remove(host_event) + return self._host_capabilities.get(host) + + def _async_discovered_by_ssdp(self, response): + @callback + def _async_start_flow(*_): + asyncio.create_task( + self._hass.config_entries.flow.async_init( + DOMAIN, + context={"source": config_entries.SOURCE_SSDP}, + data=response, + ) + ) + + # Delay starting the flow in case the discovery is the result + # of another discovery + async_call_later(self._hass, 1, _async_start_flow) + + async def _async_process_entry(self, response): + """Process a discovery.""" + _LOGGER.debug("Discovered via SSDP: %s", response) + unique_id = response["id"] + host = urlparse(response["location"]).hostname + if unique_id not in self._unique_id_capabilities: + _LOGGER.debug("Yeelight discovered with %s", response) + self._async_discovered_by_ssdp(response) + self._host_capabilities[host] = response + self._unique_id_capabilities[unique_id] = response + for event in self._host_discovered_events.get(host, []): + event.set() + if unique_id in self._callbacks: + self._hass.async_create_task(self._callbacks[unique_id](response)) + self._callbacks.pop(unique_id) + if not self._callbacks: + self._async_stop_scan() + + async def _async_start_scan(self): """Start scanning for Yeelight devices.""" _LOGGER.debug("Start scanning") - # Use loop directly to avoid home assistant track this task - self._scan_task = self._hass.loop.create_task(self._async_scan()) + await self.async_setup() + if not self._track_interval: + self._track_interval = async_track_time_interval( + self._hass, self.async_scan, DISCOVERY_INTERVAL + ) + self.async_scan() @callback def _async_stop_scan(self): """Stop scanning.""" - _LOGGER.debug("Stop scanning") - if self._scan_task is not None: - self._scan_task.cancel() - self._scan_task = None + if self._track_interval is None: + return + _LOGGER.debug("Stop scanning interval") + self._track_interval() + self._track_interval = None - @callback - def async_register_callback(self, unique_id, callback_func): + async def async_register_callback(self, unique_id, callback_func): """Register callback function.""" - host = self._seen.get(unique_id) - if host is not None: - self._hass.async_create_task(callback_func(host)) - else: - self._callbacks[unique_id] = callback_func - if len(self._callbacks) == 1: - self._async_start_scan() + if capabilities := self._unique_id_capabilities.get(unique_id): + self._hass.async_create_task(callback_func(capabilities)) + return + self._callbacks[unique_id] = callback_func + await self._async_start_scan() @callback def async_unregister_callback(self, unique_id): """Unregister callback function.""" - if unique_id not in self._callbacks: - return - self._callbacks.pop(unique_id) - if len(self._callbacks) == 0: + self._callbacks.pop(unique_id, None) + if not self._callbacks: self._async_stop_scan() class YeelightDevice: """Represents single Yeelight device.""" - def __init__(self, hass, host, config, bulb, capabilities): + def __init__(self, hass, host, config, bulb): """Initialize device.""" self._hass = hass self._config = config self._host = host self._bulb_device = bulb - self._capabilities = capabilities or {} + self._capabilities = {} self._device_type = None self._available = False self._initialized = False - - self._name = host # Default name is host - if capabilities: - # Generate name from model and id when capabilities is available - self._name = _async_unique_name(capabilities) - if config.get(CONF_NAME): - # Override default name when name is set in config - self._name = config[CONF_NAME] + self._name = None @property def bulb(self): @@ -444,7 +507,7 @@ class YeelightDevice: @property def model(self): """Return configured/autodetected device model.""" - return self._bulb_device.model + return self._bulb_device.model or self._capabilities.get("model") @property def fw_version(self): @@ -530,7 +593,8 @@ class YeelightDevice: await self.bulb.async_get_properties(UPDATE_REQUEST_PROPERTIES) self._available = True if not self._initialized: - await self._async_initialize_device() + self._initialized = True + async_dispatcher_send(self._hass, DEVICE_INITIALIZED.format(self._host)) except BulbException as ex: if self._available: # just inform once _LOGGER.error( @@ -540,28 +604,18 @@ class YeelightDevice: return self._available - async def _async_get_capabilities(self): - """Request device capabilities.""" - try: - await self._hass.async_add_executor_job(self.bulb.get_capabilities) - _LOGGER.debug( - "Device %s, %s capabilities: %s", - self._host, - self.name, - self.bulb.capabilities, - ) - except BulbException as ex: - _LOGGER.error( - "Unable to get device capabilities %s, %s: %s", - self._host, - self.name, - ex, - ) - - async def _async_initialize_device(self): - await self._async_get_capabilities() - self._initialized = True - async_dispatcher_send(self._hass, DEVICE_INITIALIZED.format(self._host)) + async def async_setup(self): + """Fetch capabilities and setup name if available.""" + scanner = YeelightScanner.async_get(self._hass) + self._capabilities = await scanner.async_get_capabilities(self._host) or {} + if name := self._config.get(CONF_NAME): + # Override default name when name is set in config + self._name = name + elif self._capabilities: + # Generate name from model and id when capabilities is available + self._name = _async_unique_name(self._capabilities) + else: + self._name = self._host # Default name is host async def async_update(self): """Update device properties and send data updated signal.""" @@ -628,6 +682,19 @@ async def _async_get_device( # Set up device bulb = AsyncBulb(host, model=model or None) - capabilities = await hass.async_add_executor_job(bulb.get_capabilities) - return YeelightDevice(hass, host, entry.options, bulb, capabilities) + device = YeelightDevice(hass, host, entry.options, bulb) + # start listening for local pushes + await device.bulb.async_listen(device.async_update_callback) + + # register stop callback to shutdown listening for local pushes + async def async_stop_listen_task(event): + """Stop listen thread.""" + _LOGGER.debug("Shutting down Yeelight Listener") + await device.bulb.async_stop_listening() + + entry.async_on_unload( + hass.bus.async_listen_once(EVENT_HOMEASSISTANT_STOP, async_stop_listen_task) + ) + + return device diff --git a/homeassistant/components/yeelight/config_flow.py b/homeassistant/components/yeelight/config_flow.py index a66571cae93..d93f59535cf 100644 --- a/homeassistant/components/yeelight/config_flow.py +++ b/homeassistant/components/yeelight/config_flow.py @@ -1,8 +1,10 @@ """Config flow for Yeelight integration.""" import logging +from urllib.parse import urlparse import voluptuous as vol import yeelight +from yeelight.aio import AsyncBulb from homeassistant import config_entries, exceptions from homeassistant.components.dhcp import IP_ADDRESS @@ -19,6 +21,7 @@ from . import ( CONF_TRANSITION, DOMAIN, NIGHTLIGHT_SWITCH_TYPE_LIGHT, + YeelightScanner, _async_unique_name, ) @@ -54,6 +57,15 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self._discovered_ip = discovery_info[IP_ADDRESS] return await self._async_handle_discovery() + async def async_step_ssdp(self, discovery_info): + """Handle discovery from ssdp.""" + self._discovered_ip = urlparse(discovery_info["location"]).hostname + await self.async_set_unique_id(discovery_info["id"]) + self._abort_if_unique_id_configured( + updates={CONF_HOST: self._discovered_ip}, reload_on_update=False + ) + return await self._async_handle_discovery() + async def _async_handle_discovery(self): """Handle any discovery.""" self.context[CONF_HOST] = self._discovered_ip @@ -62,7 +74,9 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): return self.async_abort(reason="already_in_progress") try: - self._discovered_model = await self._async_try_connect(self._discovered_ip) + self._discovered_model = await self._async_try_connect( + self._discovered_ip, raise_on_progress=True + ) except CannotConnect: return self.async_abort(reason="cannot_connect") @@ -96,7 +110,9 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): if not user_input.get(CONF_HOST): return await self.async_step_pick_device() try: - model = await self._async_try_connect(user_input[CONF_HOST]) + model = await self._async_try_connect( + user_input[CONF_HOST], raise_on_progress=False + ) except CannotConnect: errors["base"] = "cannot_connect" else: @@ -119,10 +135,12 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): if user_input is not None: unique_id = user_input[CONF_DEVICE] capabilities = self._discovered_devices[unique_id] - await self.async_set_unique_id(unique_id) + await self.async_set_unique_id(unique_id, raise_on_progress=False) self._abort_if_unique_id_configured() + host = urlparse(capabilities["location"]).hostname return self.async_create_entry( - title=_async_unique_name(capabilities), data={CONF_ID: unique_id} + title=_async_unique_name(capabilities), + data={CONF_ID: unique_id, CONF_HOST: host}, ) configured_devices = { @@ -131,19 +149,18 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): if entry.data[CONF_ID] } devices_name = {} + scanner = YeelightScanner.async_get(self.hass) + devices = await scanner.async_discover() # Run 3 times as packets can get lost - for _ in range(3): - devices = await self.hass.async_add_executor_job(yeelight.discover_bulbs) - for device in devices: - capabilities = device["capabilities"] - unique_id = capabilities["id"] - if unique_id in configured_devices: - continue # ignore configured devices - model = capabilities["model"] - host = device["ip"] - name = f"{host} {model} {unique_id}" - self._discovered_devices[unique_id] = capabilities - devices_name[unique_id] = name + for capabilities in devices: + unique_id = capabilities["id"] + if unique_id in configured_devices: + continue # ignore configured devices + model = capabilities["model"] + host = urlparse(capabilities["location"]).hostname + name = f"{host} {model} {unique_id}" + self._discovered_devices[unique_id] = capabilities + devices_name[unique_id] = name # Check if there is at least one device if not devices_name: @@ -157,7 +174,7 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): """Handle import step.""" host = user_input[CONF_HOST] try: - await self._async_try_connect(host) + await self._async_try_connect(host, raise_on_progress=False) except CannotConnect: _LOGGER.error("Failed to import %s: cannot connect", host) return self.async_abort(reason="cannot_connect") @@ -169,27 +186,26 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): self._abort_if_unique_id_configured() return self.async_create_entry(title=user_input[CONF_NAME], data=user_input) - async def _async_try_connect(self, host): + async def _async_try_connect(self, host, raise_on_progress=True): """Set up with options.""" self._async_abort_entries_match({CONF_HOST: host}) - bulb = yeelight.Bulb(host) - try: - capabilities = await self.hass.async_add_executor_job(bulb.get_capabilities) - if capabilities is None: # timeout - _LOGGER.debug("Failed to get capabilities from %s: timeout", host) - else: - _LOGGER.debug("Get capabilities: %s", capabilities) - await self.async_set_unique_id(capabilities["id"]) - return capabilities["model"] - except OSError as err: - _LOGGER.debug("Failed to get capabilities from %s: %s", host, err) - # Ignore the error since get_capabilities uses UDP discovery packet - # which does not work in all network environments - + scanner = YeelightScanner.async_get(self.hass) + capabilities = await scanner.async_get_capabilities(host) + if capabilities is None: # timeout + _LOGGER.debug("Failed to get capabilities from %s: timeout", host) + else: + _LOGGER.debug("Get capabilities: %s", capabilities) + await self.async_set_unique_id( + capabilities["id"], raise_on_progress=raise_on_progress + ) + return capabilities["model"] # Fallback to get properties + bulb = AsyncBulb(host) try: - await self.hass.async_add_executor_job(bulb.get_properties) + await bulb.async_listen(lambda _: True) + await bulb.async_get_properties() + await bulb.async_stop_listening() except yeelight.BulbException as err: _LOGGER.error("Failed to get properties from %s: %s", host, err) raise CannotConnect from err diff --git a/homeassistant/components/yeelight/light.py b/homeassistant/components/yeelight/light.py index b714ddfaba8..4766d897909 100644 --- a/homeassistant/components/yeelight/light.py +++ b/homeassistant/components/yeelight/light.py @@ -905,7 +905,7 @@ class YeelightNightLightMode(YeelightGenericLight): @property def name(self) -> str: """Return the name of the device if any.""" - return f"{self.device.name} nightlight" + return f"{self.device.name} Nightlight" @property def icon(self): @@ -997,7 +997,7 @@ class YeelightAmbientLight(YeelightColorLightWithoutNightlightSwitch): @property def name(self) -> str: """Return the name of the device if any.""" - return f"{self.device.name} ambilight" + return f"{self.device.name} Ambilight" @property def _brightness_property(self): diff --git a/homeassistant/components/yeelight/manifest.json b/homeassistant/components/yeelight/manifest.json index 3528b096c67..4c5994b1f6e 100644 --- a/homeassistant/components/yeelight/manifest.json +++ b/homeassistant/components/yeelight/manifest.json @@ -2,7 +2,7 @@ "domain": "yeelight", "name": "Yeelight", "documentation": "https://www.home-assistant.io/integrations/yeelight", - "requirements": ["yeelight==0.7.2"], + "requirements": ["yeelight==0.7.2", "async-upnp-client==0.20.0"], "codeowners": ["@rytilahti", "@zewelor", "@shenxn", "@starkillerOG"], "config_flow": true, "iot_class": "local_push", diff --git a/requirements_all.txt b/requirements_all.txt index 021d225c362..69da5abc72d 100644 --- a/requirements_all.txt +++ b/requirements_all.txt @@ -314,6 +314,7 @@ asterisk_mbox==0.5.0 # homeassistant.components.dlna_dmr # homeassistant.components.ssdp # homeassistant.components.upnp +# homeassistant.components.yeelight async-upnp-client==0.20.0 # homeassistant.components.supla diff --git a/requirements_test_all.txt b/requirements_test_all.txt index ecb66f6bad4..e476881426c 100644 --- a/requirements_test_all.txt +++ b/requirements_test_all.txt @@ -205,6 +205,7 @@ arcam-fmj==0.7.0 # homeassistant.components.dlna_dmr # homeassistant.components.ssdp # homeassistant.components.upnp +# homeassistant.components.yeelight async-upnp-client==0.20.0 # homeassistant.components.aurora diff --git a/tests/components/yeelight/__init__.py b/tests/components/yeelight/__init__.py index 9fa864d6213..cb2936cf8e2 100644 --- a/tests/components/yeelight/__init__.py +++ b/tests/components/yeelight/__init__.py @@ -1,9 +1,13 @@ """Tests for the Yeelight integration.""" +import asyncio +from datetime import timedelta from unittest.mock import AsyncMock, MagicMock, patch +from async_upnp_client.search import SSDPListener from yeelight import BulbException, BulbType from yeelight.main import _MODEL_SPECS +from homeassistant.components import yeelight as hass_yeelight from homeassistant.components.yeelight import ( CONF_MODE_MUSIC, CONF_NIGHTLIGHT_SWITCH_TYPE, @@ -13,6 +17,7 @@ from homeassistant.components.yeelight import ( YeelightScanner, ) from homeassistant.const import CONF_DEVICES, CONF_ID, CONF_NAME +from homeassistant.core import callback IP_ADDRESS = "192.168.1.239" MODEL = "color" @@ -23,13 +28,16 @@ CAPABILITIES = { "id": ID, "model": MODEL, "fw_ver": FW_VER, + "location": f"yeelight://{IP_ADDRESS}", "support": "get_prop set_default set_power toggle set_bright start_cf stop_cf" " set_scene cron_add cron_get cron_del set_ct_abx set_rgb", "name": "", } NAME = "name" -UNIQUE_NAME = f"yeelight_{MODEL}_{ID}" +SHORT_ID = hex(int("0x000000000015243f", 16)) +UNIQUE_NAME = f"yeelight_{MODEL}_{SHORT_ID}" +UNIQUE_FRIENDLY_NAME = f"Yeelight {MODEL.title()} {SHORT_ID}" MODULE = "homeassistant.components.yeelight" MODULE_CONFIG_FLOW = f"{MODULE}.config_flow" @@ -81,8 +89,8 @@ CONFIG_ENTRY_DATA = {CONF_ID: ID} def _mocked_bulb(cannot_connect=False): bulb = MagicMock() - type(bulb).get_capabilities = MagicMock( - return_value=None if cannot_connect else CAPABILITIES + type(bulb).async_listen = AsyncMock( + side_effect=BulbException if cannot_connect else None ) type(bulb).async_get_properties = AsyncMock( side_effect=BulbException if cannot_connect else None @@ -98,7 +106,6 @@ def _mocked_bulb(cannot_connect=False): bulb.last_properties = PROPERTIES.copy() bulb.music_mode = False bulb.async_get_properties = AsyncMock() - bulb.async_listen = AsyncMock() bulb.async_stop_listening = AsyncMock() bulb.async_update = AsyncMock() bulb.async_turn_on = AsyncMock() @@ -116,12 +123,43 @@ def _mocked_bulb(cannot_connect=False): return bulb -def _patch_discovery(prefix, no_device=False): +def _patched_ssdp_listener(info, *args, **kwargs): + listener = SSDPListener(*args, **kwargs) + + async def _async_callback(*_): + await listener.async_connect_callback() + + @callback + def _async_search(*_): + if info: + asyncio.create_task(listener.async_callback(info)) + + listener.async_start = _async_callback + listener.async_search = _async_search + return listener + + +def _patch_discovery(no_device=False): YeelightScanner._scanner = None # Clear class scanner to reset hass - def _mocked_discovery(timeout=2, interface=False): - if no_device: - return [] - return [{"ip": IP_ADDRESS, "port": 55443, "capabilities": CAPABILITIES}] + def _generate_fake_ssdp_listener(*args, **kwargs): + return _patched_ssdp_listener( + None if no_device else CAPABILITIES, + *args, + **kwargs, + ) - return patch(f"{prefix}.discover_bulbs", side_effect=_mocked_discovery) + return patch( + "homeassistant.components.yeelight.SSDPListener", + new=_generate_fake_ssdp_listener, + ) + + +def _patch_discovery_interval(): + return patch.object( + hass_yeelight, "DISCOVERY_SEARCH_INTERVAL", timedelta(seconds=0) + ) + + +def _patch_discovery_timeout(): + return patch.object(hass_yeelight, "DISCOVERY_TIMEOUT", 0.0001) diff --git a/tests/components/yeelight/test_binary_sensor.py b/tests/components/yeelight/test_binary_sensor.py index 472d8de4919..350c289f5b5 100644 --- a/tests/components/yeelight/test_binary_sensor.py +++ b/tests/components/yeelight/test_binary_sensor.py @@ -6,7 +6,14 @@ from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_component from homeassistant.setup import async_setup_component -from . import MODULE, NAME, PROPERTIES, YAML_CONFIGURATION, _mocked_bulb +from . import ( + MODULE, + NAME, + PROPERTIES, + YAML_CONFIGURATION, + _mocked_bulb, + _patch_discovery, +) ENTITY_BINARY_SENSOR = f"binary_sensor.{NAME}_nightlight" @@ -14,9 +21,7 @@ ENTITY_BINARY_SENSOR = f"binary_sensor.{NAME}_nightlight" async def test_nightlight(hass: HomeAssistant): """Test nightlight sensor.""" mocked_bulb = _mocked_bulb() - with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), patch( - f"{MODULE}.config_flow.yeelight.Bulb", return_value=mocked_bulb - ): + with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): await async_setup_component(hass, DOMAIN, YAML_CONFIGURATION) await hass.async_block_till_done() diff --git a/tests/components/yeelight/test_config_flow.py b/tests/components/yeelight/test_config_flow.py index 247630ecfc3..5bbfcc9283b 100644 --- a/tests/components/yeelight/test_config_flow.py +++ b/tests/components/yeelight/test_config_flow.py @@ -1,5 +1,5 @@ """Test the Yeelight config flow.""" -from unittest.mock import MagicMock, patch +from unittest.mock import patch import pytest @@ -25,14 +25,17 @@ from homeassistant.core import HomeAssistant from homeassistant.data_entry_flow import RESULT_TYPE_ABORT, RESULT_TYPE_FORM from . import ( + CAPABILITIES, ID, IP_ADDRESS, MODULE, MODULE_CONFIG_FLOW, NAME, - UNIQUE_NAME, + UNIQUE_FRIENDLY_NAME, _mocked_bulb, _patch_discovery, + _patch_discovery_interval, + _patch_discovery_timeout, ) from tests.common import MockConfigEntry @@ -55,21 +58,23 @@ async def test_discovery(hass: HomeAssistant): assert result["step_id"] == "user" assert not result["errors"] - with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight"): + with _patch_discovery(), _patch_discovery_interval(): result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) assert result2["type"] == "form" assert result2["step_id"] == "pick_device" assert not result2["errors"] - with patch(f"{MODULE}.async_setup", return_value=True) as mock_setup, patch( + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE}.async_setup", return_value=True + ) as mock_setup, patch( f"{MODULE}.async_setup_entry", return_value=True ) as mock_setup_entry: result3 = await hass.config_entries.flow.async_configure( result["flow_id"], {CONF_DEVICE: ID} ) assert result3["type"] == "create_entry" - assert result3["title"] == UNIQUE_NAME - assert result3["data"] == {CONF_ID: ID} + assert result3["title"] == UNIQUE_FRIENDLY_NAME + assert result3["data"] == {CONF_ID: ID, CONF_HOST: IP_ADDRESS} await hass.async_block_till_done() mock_setup.assert_called_once() mock_setup_entry.assert_called_once() @@ -82,7 +87,7 @@ async def test_discovery(hass: HomeAssistant): assert result["step_id"] == "user" assert not result["errors"] - with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight"): + with _patch_discovery(), _patch_discovery_interval(): result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) assert result2["type"] == "abort" assert result2["reason"] == "no_devices_found" @@ -94,7 +99,9 @@ async def test_discovery_no_device(hass: HomeAssistant): DOMAIN, context={"source": config_entries.SOURCE_USER} ) - with _patch_discovery(f"{MODULE_CONFIG_FLOW}.yeelight", no_device=True): + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(): result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) assert result2["type"] == "abort" @@ -114,26 +121,27 @@ async def test_import(hass: HomeAssistant): # Cannot connect mocked_bulb = _mocked_bulb(cannot_connect=True) - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config ) - type(mocked_bulb).get_capabilities.assert_called_once() - type(mocked_bulb).get_properties.assert_called_once() assert result["type"] == "abort" assert result["reason"] == "cannot_connect" # Success mocked_bulb = _mocked_bulb() - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb), patch( - f"{MODULE}.async_setup", return_value=True - ) as mock_setup, patch( + with _patch_discovery(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ), patch(f"{MODULE}.async_setup", return_value=True) as mock_setup, patch( f"{MODULE}.async_setup_entry", return_value=True ) as mock_setup_entry: result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config ) - type(mocked_bulb).get_capabilities.assert_called_once() assert result["type"] == "create_entry" assert result["title"] == DEFAULT_NAME assert result["data"] == { @@ -150,7 +158,9 @@ async def test_import(hass: HomeAssistant): # Duplicate mocked_bulb = _mocked_bulb() - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_IMPORT}, data=config ) @@ -169,7 +179,11 @@ async def test_manual(hass: HomeAssistant): # Cannot connect (timeout) mocked_bulb = _mocked_bulb(cannot_connect=True) - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result2 = await hass.config_entries.flow.async_configure( result["flow_id"], {CONF_HOST: IP_ADDRESS} ) @@ -178,8 +192,11 @@ async def test_manual(hass: HomeAssistant): assert result2["errors"] == {"base": "cannot_connect"} # Cannot connect (error) - type(mocked_bulb).get_capabilities = MagicMock(side_effect=OSError) - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result3 = await hass.config_entries.flow.async_configure( result["flow_id"], {CONF_HOST: IP_ADDRESS} ) @@ -187,9 +204,11 @@ async def test_manual(hass: HomeAssistant): # Success mocked_bulb = _mocked_bulb() - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb), patch( - f"{MODULE}.async_setup", return_value=True - ), patch(f"{MODULE}.async_setup_entry", return_value=True): + with _patch_discovery(), _patch_discovery_timeout(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ), patch(f"{MODULE}.async_setup", return_value=True), patch( + f"{MODULE}.async_setup_entry", return_value=True + ): result4 = await hass.config_entries.flow.async_configure( result["flow_id"], {CONF_HOST: IP_ADDRESS} ) @@ -203,7 +222,11 @@ async def test_manual(hass: HomeAssistant): DOMAIN, context={"source": config_entries.SOURCE_USER} ) mocked_bulb = _mocked_bulb() - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result2 = await hass.config_entries.flow.async_configure( result["flow_id"], {CONF_HOST: IP_ADDRESS} ) @@ -219,7 +242,7 @@ async def test_options(hass: HomeAssistant): config_entry.add_to_hass(hass) mocked_bulb = _mocked_bulb() - with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): + with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): assert await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() @@ -241,7 +264,7 @@ async def test_options(hass: HomeAssistant): config[CONF_NIGHTLIGHT_SWITCH] = True user_input = {**config} user_input.pop(CONF_NAME) - with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): + with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): result2 = await hass.config_entries.options.async_configure( result["flow_id"], user_input ) @@ -262,15 +285,18 @@ async def test_manual_no_capabilities(hass: HomeAssistant): assert not result["errors"] mocked_bulb = _mocked_bulb() - type(mocked_bulb).get_capabilities = MagicMock(return_value=None) - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb), patch( + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ), patch( f"{MODULE}.async_setup", return_value=True - ), patch(f"{MODULE}.async_setup_entry", return_value=True): + ), patch( + f"{MODULE}.async_setup_entry", return_value=True + ): result = await hass.config_entries.flow.async_configure( result["flow_id"], {CONF_HOST: IP_ADDRESS} ) - type(mocked_bulb).get_capabilities.assert_called_once() - type(mocked_bulb).get_properties.assert_called_once() assert result["type"] == "create_entry" assert result["data"] == {CONF_HOST: IP_ADDRESS} @@ -280,39 +306,53 @@ async def test_discovered_by_homekit_and_dhcp(hass): await setup.async_setup_component(hass, "persistent_notification", {}) mocked_bulb = _mocked_bulb() - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_HOMEKIT}, - data={"host": "1.2.3.4", "properties": {"id": "aa:bb:cc:dd:ee:ff"}}, + data={"host": IP_ADDRESS, "properties": {"id": "aa:bb:cc:dd:ee:ff"}}, ) + await hass.async_block_till_done() assert result["type"] == RESULT_TYPE_FORM assert result["errors"] is None - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result2 = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_DHCP}, - data={"ip": "1.2.3.4", "macaddress": "aa:bb:cc:dd:ee:ff"}, + data={"ip": IP_ADDRESS, "macaddress": "aa:bb:cc:dd:ee:ff"}, ) + await hass.async_block_till_done() assert result2["type"] == RESULT_TYPE_ABORT assert result2["reason"] == "already_in_progress" - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result3 = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_DHCP}, - data={"ip": "1.2.3.4", "macaddress": "00:00:00:00:00:00"}, + data={"ip": IP_ADDRESS, "macaddress": "00:00:00:00:00:00"}, ) + await hass.async_block_till_done() assert result3["type"] == RESULT_TYPE_ABORT assert result3["reason"] == "already_in_progress" - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", side_effect=CannotConnect): + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", side_effect=CannotConnect + ): result3 = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_DHCP}, data={"ip": "1.2.3.5", "macaddress": "00:00:00:00:00:01"}, ) + await hass.async_block_till_done() assert result3["type"] == RESULT_TYPE_ABORT assert result3["reason"] == "cannot_connect" @@ -335,17 +375,25 @@ async def test_discovered_by_dhcp_or_homekit(hass, source, data): await setup.async_setup_component(hass, "persistent_notification", {}) mocked_bulb = _mocked_bulb() - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": source}, data=data ) + await hass.async_block_till_done() + assert result["type"] == RESULT_TYPE_FORM assert result["errors"] is None - with patch(f"{MODULE}.async_setup", return_value=True) as mock_async_setup, patch( + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE}.async_setup", return_value=True + ) as mock_async_setup, patch( f"{MODULE}.async_setup_entry", return_value=True ) as mock_async_setup_entry: result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + await hass.async_block_till_done() + assert result2["type"] == "create_entry" assert result2["data"] == {CONF_HOST: IP_ADDRESS, CONF_ID: "0x000000000015243f"} assert mock_async_setup.called @@ -370,10 +418,55 @@ async def test_discovered_by_dhcp_or_homekit_failed_to_get_id(hass, source, data await setup.async_setup_component(hass, "persistent_notification", {}) mocked_bulb = _mocked_bulb() - type(mocked_bulb).get_capabilities = MagicMock(return_value=None) - with patch(f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb): + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": source}, data=data ) assert result["type"] == RESULT_TYPE_ABORT assert result["reason"] == "cannot_connect" + + +async def test_discovered_ssdp(hass): + """Test we can setup when discovered from ssdp.""" + await setup.async_setup_component(hass, "persistent_notification", {}) + + mocked_bulb = _mocked_bulb() + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_SSDP}, data=CAPABILITIES + ) + await hass.async_block_till_done() + + assert result["type"] == RESULT_TYPE_FORM + assert result["errors"] is None + + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE}.async_setup", return_value=True + ) as mock_async_setup, patch( + f"{MODULE}.async_setup_entry", return_value=True + ) as mock_async_setup_entry: + result2 = await hass.config_entries.flow.async_configure(result["flow_id"], {}) + await hass.async_block_till_done() + + assert result2["type"] == "create_entry" + assert result2["data"] == {CONF_HOST: IP_ADDRESS, CONF_ID: "0x000000000015243f"} + assert mock_async_setup.called + assert mock_async_setup_entry.called + + mocked_bulb = _mocked_bulb() + with _patch_discovery(), _patch_discovery_interval(), patch( + f"{MODULE_CONFIG_FLOW}.AsyncBulb", return_value=mocked_bulb + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_SSDP}, data=CAPABILITIES + ) + await hass.async_block_till_done() + + assert result["type"] == RESULT_TYPE_ABORT + assert result["reason"] == "already_configured" diff --git a/tests/components/yeelight/test_init.py b/tests/components/yeelight/test_init.py index 575ad4cb594..d7f4a05b436 100644 --- a/tests/components/yeelight/test_init.py +++ b/tests/components/yeelight/test_init.py @@ -1,5 +1,6 @@ """Test Yeelight.""" -from unittest.mock import AsyncMock, MagicMock, patch +from datetime import timedelta +from unittest.mock import AsyncMock, patch from yeelight import BulbException, BulbType @@ -22,9 +23,9 @@ from homeassistant.const import ( from homeassistant.core import HomeAssistant from homeassistant.helpers import entity_registry as er from homeassistant.setup import async_setup_component +from homeassistant.util import dt as dt_util from . import ( - CAPABILITIES, CONFIG_ENTRY_DATA, ENTITY_AMBILIGHT, ENTITY_BINARY_SENSOR, @@ -34,12 +35,14 @@ from . import ( ID, IP_ADDRESS, MODULE, - MODULE_CONFIG_FLOW, + SHORT_ID, _mocked_bulb, _patch_discovery, + _patch_discovery_interval, + _patch_discovery_timeout, ) -from tests.common import MockConfigEntry +from tests.common import MockConfigEntry, async_fire_time_changed async def test_ip_changes_fallback_discovery(hass: HomeAssistant): @@ -51,19 +54,15 @@ async def test_ip_changes_fallback_discovery(hass: HomeAssistant): mocked_bulb = _mocked_bulb(True) mocked_bulb.bulb_type = BulbType.WhiteTempMood - mocked_bulb.get_capabilities = MagicMock( - side_effect=[OSError, CAPABILITIES, CAPABILITIES] - ) + mocked_bulb.async_listen = AsyncMock(side_effect=[BulbException, None, None, None]) - _discovered_devices = [{"capabilities": CAPABILITIES, "ip": IP_ADDRESS}] - with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), patch( - f"{MODULE}.discover_bulbs", return_value=_discovered_devices - ): + with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery(): assert await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() + await hass.async_block_till_done() binary_sensor_entity_id = ENTITY_BINARY_SENSOR_TEMPLATE.format( - f"yeelight_color_{ID}" + f"yeelight_color_{SHORT_ID}" ) type(mocked_bulb).async_get_properties = AsyncMock(None) @@ -77,6 +76,19 @@ async def test_ip_changes_fallback_discovery(hass: HomeAssistant): entity_registry = er.async_get(hass) assert entity_registry.async_get(binary_sensor_entity_id) is not None + with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery(): + # The discovery should update the ip address + async_fire_time_changed(hass, dt_util.utcnow() + timedelta(seconds=2)) + await hass.async_block_till_done() + assert config_entry.data[CONF_HOST] == IP_ADDRESS + + # Make sure we can still reload with the new ip right after we change it + with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery(): + await hass.config_entries.async_reload(config_entry.entry_id) + await hass.async_block_till_done() + + assert entity_registry.async_get(binary_sensor_entity_id) is not None + async def test_ip_changes_id_missing_cannot_fallback(hass: HomeAssistant): """Test Yeelight ip changes and we fallback to discovery.""" @@ -85,9 +97,7 @@ async def test_ip_changes_id_missing_cannot_fallback(hass: HomeAssistant): mocked_bulb = _mocked_bulb(True) mocked_bulb.bulb_type = BulbType.WhiteTempMood - mocked_bulb.get_capabilities = MagicMock( - side_effect=[OSError, CAPABILITIES, CAPABILITIES] - ) + mocked_bulb.async_listen = AsyncMock(side_effect=[BulbException, None, None, None]) with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): assert not await hass.config_entries.async_setup(config_entry.entry_id) @@ -102,9 +112,7 @@ async def test_setup_discovery(hass: HomeAssistant): config_entry.add_to_hass(hass) mocked_bulb = _mocked_bulb() - with _patch_discovery(MODULE), patch( - f"{MODULE}.AsyncBulb", return_value=mocked_bulb - ): + with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): assert await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() @@ -127,9 +135,7 @@ async def test_setup_import(hass: HomeAssistant): """Test import from yaml.""" mocked_bulb = _mocked_bulb() name = "yeelight" - with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), patch( - f"{MODULE_CONFIG_FLOW}.yeelight.Bulb", return_value=mocked_bulb - ): + with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery(): assert await async_setup_component( hass, DOMAIN, @@ -162,9 +168,7 @@ async def test_unique_ids_device(hass: HomeAssistant): mocked_bulb = _mocked_bulb() mocked_bulb.bulb_type = BulbType.WhiteTempMood - with _patch_discovery(MODULE), patch( - f"{MODULE}.AsyncBulb", return_value=mocked_bulb - ): + with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): assert await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() @@ -188,9 +192,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}.AsyncBulb", return_value=mocked_bulb - ): + with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): assert await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() @@ -220,30 +222,13 @@ async def test_bulb_off_while_adding_in_ha(hass: HomeAssistant): mocked_bulb = _mocked_bulb(True) mocked_bulb.bulb_type = BulbType.WhiteTempMood - with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), patch( - f"{MODULE}.config_flow.yeelight.Bulb", return_value=mocked_bulb - ): + with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb), _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(): 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(".", "_") - ) - - type(mocked_bulb).get_capabilities = MagicMock(CAPABILITIES) - type(mocked_bulb).get_properties = MagicMock(None) - - await hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][ - DATA_DEVICE - ].async_update() - hass.data[DOMAIN][DATA_CONFIG_ENTRIES][config_entry.entry_id][ - DATA_DEVICE - ].async_update_callback({}) - await hass.async_block_till_done() - await hass.async_block_till_done() - - entity_registry = er.async_get(hass) - assert entity_registry.async_get(binary_sensor_entity_id) is not None + assert config_entry.state is ConfigEntryState.LOADED async def test_async_listen_error_late_discovery(hass, caplog): @@ -251,12 +236,9 @@ async def test_async_listen_error_late_discovery(hass, caplog): config_entry = MockConfigEntry(domain=DOMAIN, data=CONFIG_ENTRY_DATA) config_entry.add_to_hass(hass) - mocked_bulb = _mocked_bulb() - mocked_bulb.async_listen = AsyncMock(side_effect=BulbException) + mocked_bulb = _mocked_bulb(cannot_connect=True) - with _patch_discovery(MODULE), patch( - f"{MODULE}.AsyncBulb", return_value=mocked_bulb - ): + with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): await hass.config_entries.async_setup(config_entry.entry_id) await hass.async_block_till_done() @@ -264,17 +246,33 @@ async def test_async_listen_error_late_discovery(hass, caplog): assert "Failed to connect to bulb at" in caplog.text -async def test_async_listen_error_has_host(hass: HomeAssistant): +async def test_async_listen_error_has_host_with_id(hass: HomeAssistant): """Test the async listen error.""" config_entry = MockConfigEntry( domain=DOMAIN, data={CONF_ID: ID, CONF_HOST: "127.0.0.1"} ) config_entry.add_to_hass(hass) - mocked_bulb = _mocked_bulb() - mocked_bulb.async_listen = AsyncMock(side_effect=BulbException) + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE}.AsyncBulb", return_value=_mocked_bulb(cannot_connect=True) + ): + await hass.config_entries.async_setup(config_entry.entry_id) - with patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): + assert config_entry.state is ConfigEntryState.LOADED + + +async def test_async_listen_error_has_host_without_id(hass: HomeAssistant): + """Test the async listen error but no id.""" + config_entry = MockConfigEntry(domain=DOMAIN, data={CONF_HOST: "127.0.0.1"}) + config_entry.add_to_hass(hass) + + with _patch_discovery( + no_device=True + ), _patch_discovery_timeout(), _patch_discovery_interval(), patch( + f"{MODULE}.AsyncBulb", return_value=_mocked_bulb(cannot_connect=True) + ): await hass.config_entries.async_setup(config_entry.entry_id) assert config_entry.state is ConfigEntryState.SETUP_RETRY diff --git a/tests/components/yeelight/test_light.py b/tests/components/yeelight/test_light.py index 8b7ec154b83..7497fa8773e 100644 --- a/tests/components/yeelight/test_light.py +++ b/tests/components/yeelight/test_light.py @@ -102,9 +102,10 @@ from . import ( MODULE, NAME, PROPERTIES, - UNIQUE_NAME, + UNIQUE_FRIENDLY_NAME, _mocked_bulb, _patch_discovery, + _patch_discovery_interval, ) from tests.common import MockConfigEntry @@ -132,7 +133,7 @@ async def test_services(hass: HomeAssistant, caplog): config_entry.add_to_hass(hass) mocked_bulb = _mocked_bulb() - with _patch_discovery(MODULE), patch( + with _patch_discovery(), _patch_discovery_interval(), patch( f"{MODULE}.AsyncBulb", return_value=mocked_bulb ): assert await hass.config_entries.async_setup(config_entry.entry_id) @@ -559,7 +560,7 @@ async def test_device_types(hass: HomeAssistant, caplog): model, target_properties, nightlight_properties=None, - name=UNIQUE_NAME, + name=UNIQUE_FRIENDLY_NAME, entity_id=ENTITY_LIGHT, ): config_entry = MockConfigEntry( @@ -598,7 +599,7 @@ async def test_device_types(hass: HomeAssistant, caplog): assert hass.states.get(entity_id).state == "off" state = hass.states.get(f"{entity_id}_nightlight") assert state.state == "on" - nightlight_properties["friendly_name"] = f"{name} nightlight" + nightlight_properties["friendly_name"] = f"{name} Nightlight" nightlight_properties["icon"] = "mdi:weather-night" nightlight_properties["flowing"] = False nightlight_properties["night_light"] = True @@ -893,7 +894,7 @@ async def test_device_types(hass: HomeAssistant, caplog): "color_mode": "color_temp", "supported_color_modes": ["color_temp", "hs", "rgb"], }, - name=f"{UNIQUE_NAME} ambilight", + name=f"{UNIQUE_FRIENDLY_NAME} Ambilight", entity_id=f"{ENTITY_LIGHT}_ambilight", ) @@ -914,7 +915,7 @@ async def test_device_types(hass: HomeAssistant, caplog): "color_mode": "hs", "supported_color_modes": ["color_temp", "hs", "rgb"], }, - name=f"{UNIQUE_NAME} ambilight", + name=f"{UNIQUE_FRIENDLY_NAME} Ambilight", entity_id=f"{ENTITY_LIGHT}_ambilight", ) @@ -935,7 +936,7 @@ async def test_device_types(hass: HomeAssistant, caplog): "color_mode": "rgb", "supported_color_modes": ["color_temp", "hs", "rgb"], }, - name=f"{UNIQUE_NAME} ambilight", + name=f"{UNIQUE_FRIENDLY_NAME} Ambilight", entity_id=f"{ENTITY_LIGHT}_ambilight", ) @@ -969,7 +970,7 @@ async def test_effects(hass: HomeAssistant): config_entry.add_to_hass(hass) mocked_bulb = _mocked_bulb() - with _patch_discovery(MODULE), patch( + with _patch_discovery(), _patch_discovery_interval(), patch( f"{MODULE}.AsyncBulb", return_value=mocked_bulb ): assert await hass.config_entries.async_setup(config_entry.entry_id)