From c3c0a5292918ba73cf595bad139ff67171b54cc6 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Tue, 26 Oct 2021 14:04:19 -0500 Subject: [PATCH] Fix overriding the yeelight model if it is not known (#56967) --- homeassistant/components/yeelight/__init__.py | 42 ++++++------- .../components/yeelight/config_flow.py | 60 ++++++++++++------- .../components/yeelight/strings.json | 3 +- .../components/yeelight/translations/en.json | 5 +- tests/components/yeelight/test_config_flow.py | 47 ++++++++++++++- tests/components/yeelight/test_init.py | 4 +- 6 files changed, 110 insertions(+), 51 deletions(-) diff --git a/homeassistant/components/yeelight/__init__.py b/homeassistant/components/yeelight/__init__.py index a1463daed12..04b3bc4d0a8 100644 --- a/homeassistant/components/yeelight/__init__.py +++ b/homeassistant/components/yeelight/__init__.py @@ -46,6 +46,7 @@ DEFAULT_SAVE_ON_CHANGE = False DEFAULT_NIGHTLIGHT_SWITCH = False CONF_MODEL = "model" +CONF_DETECTED_MODEL = "detected_model" CONF_TRANSITION = "transition" CONF_SAVE_ON_CHANGE = "save_on_change" CONF_MODE_MUSIC = "use_music_mode" @@ -203,15 +204,15 @@ async def _async_initialize( if ( device.capabilities - and entry.options.get(CONF_MODEL) != device.capabilities["model"] + and entry.data.get(CONF_DETECTED_MODEL) != device.capabilities["model"] ): hass.config_entries.async_update_entry( - entry, options={**entry.options, CONF_MODEL: device.capabilities["model"]} + entry, + data={**entry.data, CONF_DETECTED_MODEL: device.capabilities["model"]}, ) # fetch initial state await device.async_update() - entry.async_on_unload(entry.add_update_listener(_async_update_listener)) @callback @@ -228,10 +229,13 @@ def _async_normalize_config_entry(hass: HomeAssistant, entry: ConfigEntry) -> No data={ CONF_HOST: entry.data.get(CONF_HOST), CONF_ID: entry.data.get(CONF_ID) or entry.unique_id, + CONF_DETECTED_MODEL: entry.data.get(CONF_DETECTED_MODEL), }, options={ CONF_NAME: entry.data.get(CONF_NAME, ""), - CONF_MODEL: entry.data.get(CONF_MODEL, ""), + CONF_MODEL: entry.data.get( + CONF_MODEL, entry.data.get(CONF_DETECTED_MODEL, "") + ), CONF_TRANSITION: entry.data.get(CONF_TRANSITION, DEFAULT_TRANSITION), CONF_MODE_MUSIC: entry.data.get(CONF_MODE_MUSIC, DEFAULT_MODE_MUSIC), CONF_SAVE_ON_CHANGE: entry.data.get( @@ -271,6 +275,9 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool: hass.config_entries.async_setup_platforms(entry, PLATFORMS) + # Wait to install the reload listener until everything was successfully initialized + entry.async_on_unload(entry.add_update_listener(_async_update_listener)) + return True @@ -658,37 +665,30 @@ class YeelightDevice: class YeelightEntity(Entity): """Represents single Yeelight entity.""" + _attr_should_poll = False + def __init__(self, device: YeelightDevice, entry: ConfigEntry) -> None: """Initialize the entity.""" self._device = device self._unique_id = entry.unique_id or entry.entry_id + self._attr_device_info = DeviceInfo( + identifiers={(DOMAIN, self._unique_id)}, + name=self._device.name, + manufacturer="Yeelight", + model=self._device.model, + sw_version=self._device.fw_version, + ) @property def unique_id(self) -> str: """Return the unique ID.""" return self._unique_id - @property - def device_info(self) -> DeviceInfo: - """Return the device info.""" - return { - "identifiers": {(DOMAIN, self._unique_id)}, - "name": self._device.name, - "manufacturer": "Yeelight", - "model": self._device.model, - "sw_version": self._device.fw_version, - } - @property def available(self) -> bool: """Return if bulb is available.""" return self._device.available - @property - def should_poll(self) -> bool: - """No polling needed.""" - return False - async def async_update(self) -> None: """Update the entity.""" await self._device.async_update() @@ -698,7 +698,7 @@ async def _async_get_device( hass: HomeAssistant, host: str, entry: ConfigEntry ) -> YeelightDevice: # Get model from config and capabilities - model = entry.options.get(CONF_MODEL) + model = entry.options.get(CONF_MODEL) or entry.data.get(CONF_DETECTED_MODEL) # Set up device bulb = AsyncBulb(host, model=model or None) diff --git a/homeassistant/components/yeelight/config_flow.py b/homeassistant/components/yeelight/config_flow.py index f4aab12a34a..fbc9270d72a 100644 --- a/homeassistant/components/yeelight/config_flow.py +++ b/homeassistant/components/yeelight/config_flow.py @@ -5,6 +5,7 @@ from urllib.parse import urlparse import voluptuous as vol import yeelight from yeelight.aio import AsyncBulb +from yeelight.main import get_known_models from homeassistant import config_entries, exceptions from homeassistant.components.dhcp import IP_ADDRESS @@ -14,6 +15,7 @@ from homeassistant.core import callback import homeassistant.helpers.config_validation as cv from . import ( + CONF_DETECTED_MODEL, CONF_MODE_MUSIC, CONF_MODEL, CONF_NIGHTLIGHT_SWITCH, @@ -268,31 +270,45 @@ class OptionsFlowHandler(config_entries.OptionsFlow): async def async_step_init(self, user_input=None): """Handle the initial step.""" - if user_input is not None: - options = {**self._config_entry.options} - options.update(user_input) - return self.async_create_entry(title="", data=options) - + data = self._config_entry.data options = self._config_entry.options + detected_model = data.get(CONF_DETECTED_MODEL) + model = options[CONF_MODEL] or detected_model + + if user_input is not None: + return self.async_create_entry( + title="", data={CONF_MODEL: model, **options, **user_input} + ) + + schema_dict = {} + known_models = get_known_models() + if is_unknown_model := model not in known_models: + known_models.insert(0, model) + + if is_unknown_model or model != detected_model: + schema_dict.update( + { + vol.Optional(CONF_MODEL, default=model): vol.In(known_models), + } + ) + schema_dict.update( + { + vol.Required( + CONF_TRANSITION, default=options[CONF_TRANSITION] + ): cv.positive_int, + vol.Required(CONF_MODE_MUSIC, default=options[CONF_MODE_MUSIC]): bool, + vol.Required( + CONF_SAVE_ON_CHANGE, default=options[CONF_SAVE_ON_CHANGE] + ): bool, + vol.Required( + CONF_NIGHTLIGHT_SWITCH, default=options[CONF_NIGHTLIGHT_SWITCH] + ): bool, + } + ) + return self.async_show_form( step_id="init", - data_schema=vol.Schema( - { - vol.Optional(CONF_MODEL, default=options[CONF_MODEL]): str, - vol.Required( - CONF_TRANSITION, default=options[CONF_TRANSITION] - ): cv.positive_int, - vol.Required( - CONF_MODE_MUSIC, default=options[CONF_MODE_MUSIC] - ): bool, - vol.Required( - CONF_SAVE_ON_CHANGE, default=options[CONF_SAVE_ON_CHANGE] - ): bool, - vol.Required( - CONF_NIGHTLIGHT_SWITCH, default=options[CONF_NIGHTLIGHT_SWITCH] - ): bool, - } - ), + data_schema=vol.Schema(schema_dict), ) diff --git a/homeassistant/components/yeelight/strings.json b/homeassistant/components/yeelight/strings.json index 73868b6c571..2614b7b7899 100644 --- a/homeassistant/components/yeelight/strings.json +++ b/homeassistant/components/yeelight/strings.json @@ -28,9 +28,8 @@ "options": { "step": { "init": { - "description": "If you leave model empty, it will be automatically detected.", "data": { - "model": "Model (optional)", + "model": "Model", "transition": "Transition Time (ms)", "use_music_mode": "Enable Music Mode", "save_on_change": "Save Status On Change", diff --git a/homeassistant/components/yeelight/translations/en.json b/homeassistant/components/yeelight/translations/en.json index 4ed9440aa8f..b49e3d7658e 100644 --- a/homeassistant/components/yeelight/translations/en.json +++ b/homeassistant/components/yeelight/translations/en.json @@ -29,13 +29,12 @@ "step": { "init": { "data": { - "model": "Model (optional)", + "model": "Model", "nightlight_switch": "Use Nightlight Switch", "save_on_change": "Save Status On Change", "transition": "Transition Time (ms)", "use_music_mode": "Enable Music Mode" - }, - "description": "If you leave model empty, it will be automatically detected." + } } } } diff --git a/tests/components/yeelight/test_config_flow.py b/tests/components/yeelight/test_config_flow.py index 85d147ba697..258477c9569 100644 --- a/tests/components/yeelight/test_config_flow.py +++ b/tests/components/yeelight/test_config_flow.py @@ -5,6 +5,7 @@ import pytest from homeassistant import config_entries from homeassistant.components.yeelight import ( + CONF_DETECTED_MODEL, CONF_MODE_MUSIC, CONF_MODEL, CONF_NIGHTLIGHT_SWITCH, @@ -332,7 +333,8 @@ async def test_manual(hass: HomeAssistant): async def test_options(hass: HomeAssistant): """Test options flow.""" config_entry = MockConfigEntry( - domain=DOMAIN, data={CONF_HOST: IP_ADDRESS, CONF_NAME: NAME} + domain=DOMAIN, + data={CONF_HOST: IP_ADDRESS, CONF_NAME: NAME, CONF_DETECTED_MODEL: MODEL}, ) config_entry.add_to_hass(hass) @@ -356,6 +358,49 @@ async def test_options(hass: HomeAssistant): assert result["type"] == "form" assert result["step_id"] == "init" + config[CONF_NIGHTLIGHT_SWITCH] = True + user_input = {**config} + user_input.pop(CONF_NAME) + user_input.pop(CONF_MODEL) + with _patch_discovery(), patch(f"{MODULE}.AsyncBulb", return_value=mocked_bulb): + result2 = await hass.config_entries.options.async_configure( + result["flow_id"], user_input + ) + await hass.async_block_till_done() + assert result2["type"] == "create_entry" + assert result2["data"] == config + assert result2["data"] == config_entry.options + assert hass.states.get(f"light.{NAME}_nightlight") is not None + + +async def test_options_unknown_model(hass: HomeAssistant): + """Test options flow with an unknown model.""" + config_entry = MockConfigEntry( + domain=DOMAIN, + data={CONF_HOST: IP_ADDRESS, CONF_NAME: NAME, CONF_DETECTED_MODEL: "not_in_db"}, + ) + config_entry.add_to_hass(hass) + + mocked_bulb = _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() + + config = { + CONF_NAME: NAME, + CONF_MODEL: "not_in_db", + CONF_TRANSITION: DEFAULT_TRANSITION, + CONF_MODE_MUSIC: DEFAULT_MODE_MUSIC, + CONF_SAVE_ON_CHANGE: DEFAULT_SAVE_ON_CHANGE, + CONF_NIGHTLIGHT_SWITCH: DEFAULT_NIGHTLIGHT_SWITCH, + } + assert config_entry.options == config + assert hass.states.get(f"light.{NAME}_nightlight") is None + + result = await hass.config_entries.options.async_init(config_entry.entry_id) + assert result["type"] == "form" + assert result["step_id"] == "init" + config[CONF_NIGHTLIGHT_SWITCH] = True user_input = {**config} user_input.pop(CONF_NAME) diff --git a/tests/components/yeelight/test_init.py b/tests/components/yeelight/test_init.py index 73d2543f9d4..16d8d7c8a90 100644 --- a/tests/components/yeelight/test_init.py +++ b/tests/components/yeelight/test_init.py @@ -8,7 +8,7 @@ from yeelight import BulbException, BulbType from yeelight.aio import KEY_CONNECTED from homeassistant.components.yeelight import ( - CONF_MODEL, + CONF_DETECTED_MODEL, CONF_NIGHTLIGHT_SWITCH, CONF_NIGHTLIGHT_SWITCH_TYPE, DOMAIN, @@ -377,7 +377,7 @@ async def test_async_listen_error_late_discovery(hass, caplog): await hass.async_block_till_done() assert config_entry.state is ConfigEntryState.LOADED - assert config_entry.options[CONF_MODEL] == MODEL + assert config_entry.data[CONF_DETECTED_MODEL] == MODEL async def test_unload_before_discovery(hass, caplog):