From 963651d6f23fd4903e862535956f64b1f90ac0d2 Mon Sep 17 00:00:00 2001 From: On Freund Date: Tue, 1 Sep 2020 15:08:37 +0300 Subject: [PATCH] Add support for authenticated Shelly devices (#39461) * Add support for authenticated Shelly devices * Fix comment typos * Update homeassistant/components/shelly/config_flow.py Co-authored-by: Maciej Bieniek * Fix unauthenticated devices * Update homeassistant/components/shelly/config_flow.py Co-authored-by: Martin Hjelmare * Code review fixes * More code review fixes * Fix typo * Update homeassistant/components/shelly/config_flow.py Co-authored-by: Martin Hjelmare Co-authored-by: Maciej Bieniek Co-authored-by: Martin Hjelmare --- homeassistant/components/shelly/__init__.py | 12 +- .../components/shelly/config_flow.py | 71 +++++++-- homeassistant/components/shelly/strings.json | 10 +- .../components/shelly/translations/en.json | 16 +- tests/components/shelly/test_config_flow.py | 142 +++++++++++++++++- 5 files changed, 220 insertions(+), 31 deletions(-) diff --git a/homeassistant/components/shelly/__init__.py b/homeassistant/components/shelly/__init__.py index 427bd0148f2..ebc3d287e5e 100644 --- a/homeassistant/components/shelly/__init__.py +++ b/homeassistant/components/shelly/__init__.py @@ -8,7 +8,12 @@ import aioshelly import async_timeout from homeassistant.config_entries import ConfigEntry -from homeassistant.const import EVENT_HOMEASSISTANT_STOP +from homeassistant.const import ( + CONF_HOST, + CONF_PASSWORD, + CONF_USERNAME, + EVENT_HOMEASSISTANT_STOP, +) from homeassistant.core import HomeAssistant, callback from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import ( @@ -35,7 +40,10 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry): try: async with async_timeout.timeout(5): device = await aioshelly.Device.create( - entry.data["host"], aiohttp_client.async_get_clientsession(hass) + entry.data[CONF_HOST], + aiohttp_client.async_get_clientsession(hass), + entry.data.get(CONF_USERNAME), + entry.data.get(CONF_PASSWORD), ) except (asyncio.TimeoutError, OSError) as err: raise ConfigEntryNotReady from err diff --git a/homeassistant/components/shelly/config_flow.py b/homeassistant/components/shelly/config_flow.py index c464f0d7adb..cd35f7d2552 100644 --- a/homeassistant/components/shelly/config_flow.py +++ b/homeassistant/components/shelly/config_flow.py @@ -8,25 +8,29 @@ import async_timeout import voluptuous as vol from homeassistant import config_entries, core +from homeassistant.const import CONF_HOST, CONF_PASSWORD, CONF_USERNAME from homeassistant.helpers import aiohttp_client from .const import DOMAIN # pylint:disable=unused-import _LOGGER = logging.getLogger(__name__) -DATA_SCHEMA = vol.Schema({"host": str}) +HOST_SCHEMA = vol.Schema({vol.Required(CONF_HOST): str}) HTTP_CONNECT_ERRORS = (asyncio.TimeoutError, aiohttp.ClientError) -async def validate_input(hass: core.HomeAssistant, data): +async def validate_input(hass: core.HomeAssistant, host, data): """Validate the user input allows us to connect. Data has the keys from DATA_SCHEMA with values provided by the user. """ async with async_timeout.timeout(5): device = await aioshelly.Device.create( - data["host"], aiohttp_client.async_get_clientsession(hass) + host, + aiohttp_client.async_get_clientsession(hass), + data.get(CONF_USERNAME), + data.get(CONF_PASSWORD), ) await device.shutdown() @@ -47,34 +51,71 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): """Handle the initial step.""" errors = {} if user_input is not None: + host = user_input[CONF_HOST] try: - info = await self._async_get_info(user_input["host"]) + info = await self._async_get_info(host) except HTTP_CONNECT_ERRORS: errors["base"] = "cannot_connect" except Exception: # pylint: disable=broad-except _LOGGER.exception("Unexpected exception") errors["base"] = "unknown" - else: + await self.async_set_unique_id(info["mac"]) + self._abort_if_unique_id_configured({CONF_HOST: host}) + self.host = host if info["auth"]: - return self.async_abort(reason="auth_not_supported") + return await self.async_step_credentials() try: - device_info = await validate_input(self.hass, user_input) + device_info = await validate_input(self.hass, self.host, {}) except HTTP_CONNECT_ERRORS: errors["base"] = "cannot_connect" except Exception: # pylint: disable=broad-except _LOGGER.exception("Unexpected exception") errors["base"] = "unknown" else: - await self.async_set_unique_id(device_info["mac"]) return self.async_create_entry( - title=device_info["title"] or user_input["host"], + title=device_info["title"] or self.host, data=user_input, ) return self.async_show_form( - step_id="user", data_schema=DATA_SCHEMA, errors=errors + step_id="user", data_schema=HOST_SCHEMA, errors=errors + ) + + async def async_step_credentials(self, user_input=None): + """Handle the credentials step.""" + errors = {} + if user_input is not None: + try: + device_info = await validate_input(self.hass, self.host, user_input) + except aiohttp.ClientResponseError as error: + if error.status == 401: + errors["base"] = "invalid_auth" + else: + errors["base"] = "cannot_connect" + except HTTP_CONNECT_ERRORS: + errors["base"] = "cannot_connect" + except Exception: # pylint: disable=broad-except + _LOGGER.exception("Unexpected exception") + errors["base"] = "unknown" + else: + return self.async_create_entry( + title=device_info["title"] or self.host, + data={**user_input, CONF_HOST: self.host}, + ) + else: + user_input = {} + + schema = vol.Schema( + { + vol.Required(CONF_USERNAME, default=user_input.get(CONF_USERNAME)): str, + vol.Required(CONF_PASSWORD, default=user_input.get(CONF_PASSWORD)): str, + } + ) + + return self.async_show_form( + step_id="credentials", data_schema=schema, errors=errors ) async def async_step_zeroconf(self, zeroconf_info): @@ -87,11 +128,8 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): except HTTP_CONNECT_ERRORS: return self.async_abort(reason="cannot_connect") - if info["auth"]: - return self.async_abort(reason="auth_not_supported") - await self.async_set_unique_id(info["mac"]) - self._abort_if_unique_id_configured({"host": zeroconf_info["host"]}) + self._abort_if_unique_id_configured({CONF_HOST: zeroconf_info["host"]}) self.host = zeroconf_info["host"] # pylint: disable=no-member # https://github.com/PyCQA/pylint/issues/3167 self.context["title_placeholders"] = {"name": zeroconf_info["host"]} @@ -101,8 +139,11 @@ class ConfigFlow(config_entries.ConfigFlow, domain=DOMAIN): """Handle discovery confirm.""" errors = {} if user_input is not None: + if self.info["auth"]: + return await self.async_step_credentials() + try: - device_info = await validate_input(self.hass, {"host": self.host}) + device_info = await validate_input(self.hass, self.host, {}) except HTTP_CONNECT_ERRORS: errors["base"] = "cannot_connect" except Exception: # pylint: disable=broad-except diff --git a/homeassistant/components/shelly/strings.json b/homeassistant/components/shelly/strings.json index 9c5f5707914..16dc331e452 100644 --- a/homeassistant/components/shelly/strings.json +++ b/homeassistant/components/shelly/strings.json @@ -8,14 +8,20 @@ "host": "[%key:common::config_flow::data::host%]" } }, + "credentials": { + "data": { + "username": "[%key:common::config_flow::data::username%]", + "password": "[%key:common::config_flow::data::password%]" + } + }, "confirm_discovery": { "description": "Do you want to set up the {model} at {host}?" } }, "error": { "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]", - "unknown": "[%key:common::config_flow::error::unknown%]", - "auth_not_supported": "Shelly devices requiring authentication are not currently supported." + "invalid_auth": "[%key:common::config_flow::error::invalid_auth%]", + "unknown": "[%key:common::config_flow::error::unknown%]" }, "abort": { "already_configured": "[%key:common::config_flow::abort::already_configured_device%]" diff --git a/homeassistant/components/shelly/translations/en.json b/homeassistant/components/shelly/translations/en.json index ebcb1976516..5c3e86e73ad 100644 --- a/homeassistant/components/shelly/translations/en.json +++ b/homeassistant/components/shelly/translations/en.json @@ -1,21 +1,27 @@ { "config": { "abort": { - "already_configured": "Device is already configured" + "already_configured": "[%key:common::config_flow::abort::already_configured_device%]" }, "error": { - "auth_not_supported": "Shelly devices requiring authentication are not currently supported.", - "cannot_connect": "Failed to connect", - "unknown": "Unexpected error" + "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]", + "invalid_auth": "[%key:common::config_flow::error::invalid_auth%]", + "unknown": "[%key:common::config_flow::error::unknown%]" }, "flow_title": "Shelly: {name}", "step": { "confirm_discovery": { "description": "Do you want to set up the {model} at {host}?" }, + "credentials": { + "data": { + "password": "[%key:common::config_flow::data::password%]", + "username": "[%key:common::config_flow::data::username%]" + } + }, "user": { "data": { - "host": "Host" + "host": "[%key:common::config_flow::data::host%]" } } } diff --git a/tests/components/shelly/test_config_flow.py b/tests/components/shelly/test_config_flow.py index 9acd52b6e8e..93192e89df3 100644 --- a/tests/components/shelly/test_config_flow.py +++ b/tests/components/shelly/test_config_flow.py @@ -1,6 +1,7 @@ """Test the Shelly config flow.""" import asyncio +import aiohttp import pytest from homeassistant import config_entries, setup @@ -50,7 +51,7 @@ async def test_form(hass): async def test_form_auth(hass): - """Test we can't manually configure if auth is required.""" + """Test manual configuration if auth is required.""" result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": config_entries.SOURCE_USER} ) @@ -66,8 +67,36 @@ async def test_form_auth(hass): {"host": "1.1.1.1"}, ) - assert result2["type"] == "abort" - assert result2["reason"] == "auth_not_supported" + assert result2["type"] == "form" + assert result["errors"] == {} + + with patch( + "aioshelly.Device.create", + return_value=Mock( + shutdown=AsyncMock(), + settings={"name": "Test name", "device": {"mac": "test-mac"}}, + ), + ), patch( + "homeassistant.components.shelly.async_setup", return_value=True + ) as mock_setup, patch( + "homeassistant.components.shelly.async_setup_entry", + return_value=True, + ) as mock_setup_entry: + result3 = await hass.config_entries.flow.async_configure( + result2["flow_id"], + {"username": "test username", "password": "test password"}, + ) + + assert result3["type"] == "create_entry" + assert result3["title"] == "Test name" + assert result3["data"] == { + "host": "1.1.1.1", + "username": "test username", + "password": "test password", + } + await hass.async_block_till_done() + assert len(mock_setup.mock_calls) == 1 + assert len(mock_setup_entry.mock_calls) == 1 @pytest.mark.parametrize( @@ -103,7 +132,9 @@ async def test_form_errors_test_connection(hass, error): DOMAIN, context={"source": config_entries.SOURCE_USER} ) - with patch("aioshelly.get_info", return_value={"auth": False}), patch( + with patch( + "aioshelly.get_info", return_value={"mac": "test-mac", "auth": False} + ), patch( "aioshelly.Device.create", side_effect=exc, ): @@ -116,6 +147,68 @@ async def test_form_errors_test_connection(hass, error): assert result2["errors"] == {"base": base_error} +async def test_form_already_configured(hass): + """Test we get the form.""" + await setup.async_setup_component(hass, "persistent_notification", {}) + entry = MockConfigEntry( + domain="shelly", unique_id="test-mac", data={"host": "0.0.0.0"} + ) + entry.add_to_hass(hass) + + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + + with patch( + "aioshelly.get_info", + return_value={"mac": "test-mac", "type": "SHSW-1", "auth": False}, + ): + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"host": "1.1.1.1"}, + ) + + assert result2["type"] == "abort" + assert result2["reason"] == "already_configured" + + # Test config entry got updated with latest IP + assert entry.data["host"] == "1.1.1.1" + + +@pytest.mark.parametrize( + "error", + [ + (aiohttp.ClientResponseError(Mock(), (), status=400), "cannot_connect"), + (aiohttp.ClientResponseError(Mock(), (), status=401), "invalid_auth"), + (asyncio.TimeoutError, "cannot_connect"), + (ValueError, "unknown"), + ], +) +async def test_form_auth_errors_test_connection(hass, error): + """Test we handle errors in authenticated devices.""" + exc, base_error = error + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": config_entries.SOURCE_USER} + ) + + with patch("aioshelly.get_info", return_value={"mac": "test-mac", "auth": True}): + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], + {"host": "1.1.1.1"}, + ) + + with patch( + "aioshelly.Device.create", + side_effect=exc, + ): + result3 = await hass.config_entries.flow.async_configure( + result2["flow_id"], + {"username": "test username", "password": "test password"}, + ) + assert result3["type"] == "form" + assert result3["errors"] == {"base": base_error} + + async def test_zeroconf(hass): """Test we get the form.""" await setup.async_setup_component(hass, "persistent_notification", {}) @@ -232,7 +325,7 @@ async def test_zeroconf_cannot_connect(hass): async def test_zeroconf_require_auth(hass): - """Test we get the form.""" + """Test zeroconf if auth is required.""" await setup.async_setup_component(hass, "persistent_notification", {}) with patch( @@ -244,8 +337,43 @@ async def test_zeroconf_require_auth(hass): data={"host": "1.1.1.1", "name": "shelly1pm-12345"}, context={"source": config_entries.SOURCE_ZEROCONF}, ) - assert result["type"] == "abort" - assert result["reason"] == "auth_not_supported" + assert result["type"] == "form" + assert result["errors"] == {} + + result2 = await hass.config_entries.flow.async_configure( + result["flow_id"], + {}, + ) + assert result2["type"] == "form" + assert result2["errors"] == {} + + with patch( + "aioshelly.Device.create", + return_value=Mock( + shutdown=AsyncMock(), + settings={"name": "Test name", "device": {"mac": "test-mac"}}, + ), + ), patch( + "homeassistant.components.shelly.async_setup", return_value=True + ) as mock_setup, patch( + "homeassistant.components.shelly.async_setup_entry", + return_value=True, + ) as mock_setup_entry: + result3 = await hass.config_entries.flow.async_configure( + result2["flow_id"], + {"username": "test username", "password": "test password"}, + ) + + assert result3["type"] == "create_entry" + assert result3["title"] == "Test name" + assert result3["data"] == { + "host": "1.1.1.1", + "username": "test username", + "password": "test password", + } + await hass.async_block_till_done() + assert len(mock_setup.mock_calls) == 1 + assert len(mock_setup_entry.mock_calls) == 1 async def test_zeroconf_not_shelly(hass):