From f4561891ae59d378f5940e4ed332a6a6e99b6308 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Mon, 9 Mar 2020 21:39:42 -0600 Subject: [PATCH] Ensure AirVisual import config flow checks API key correctness (#32624) * Add small improvements to AirVisual config flow tests * Code review comments * Code review comments --- .../components/airvisual/config_flow.py | 47 +++++++++---------- .../components/airvisual/test_config_flow.py | 33 ++++++------- 2 files changed, 38 insertions(+), 42 deletions(-) diff --git a/homeassistant/components/airvisual/config_flow.py b/homeassistant/components/airvisual/config_flow.py index 4dd7fb80de8..bdd1d9a7b70 100644 --- a/homeassistant/components/airvisual/config_flow.py +++ b/homeassistant/components/airvisual/config_flow.py @@ -1,4 +1,6 @@ """Define a config flow manager for AirVisual.""" +import logging + from pyairvisual import Client from pyairvisual.errors import InvalidKeyError import voluptuous as vol @@ -10,6 +12,8 @@ from homeassistant.helpers import aiohttp_client, config_validation as cv from .const import CONF_GEOGRAPHIES, DOMAIN # pylint: disable=unused-import +_LOGGER = logging.getLogger("homeassistant.components.airvisual") + class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): """Handle a AirVisual config flow.""" @@ -46,21 +50,7 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): async def async_step_import(self, import_config): """Import a config entry from configuration.yaml.""" - await self._async_set_unique_id(import_config[CONF_API_KEY]) - - data = {**import_config} - if not data.get(CONF_GEOGRAPHIES): - data[CONF_GEOGRAPHIES] = [ - { - CONF_LATITUDE: self.hass.config.latitude, - CONF_LONGITUDE: self.hass.config.longitude, - } - ] - - return self.async_create_entry( - title=f"Cloud API (API key: {import_config[CONF_API_KEY][:4]}...)", - data=data, - ) + return await self.async_step_user(import_config) async def async_step_user(self, user_input=None): """Handle the start of the config flow.""" @@ -70,7 +60,6 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): await self._async_set_unique_id(user_input[CONF_API_KEY]) websession = aiohttp_client.async_get_clientsession(self.hass) - client = Client(websession, api_key=user_input[CONF_API_KEY]) try: @@ -78,15 +67,21 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): except InvalidKeyError: return await self._show_form(errors={CONF_API_KEY: "invalid_api_key"}) + data = {CONF_API_KEY: user_input[CONF_API_KEY]} + if user_input.get(CONF_GEOGRAPHIES): + data[CONF_GEOGRAPHIES] = user_input[CONF_GEOGRAPHIES] + else: + data[CONF_GEOGRAPHIES] = [ + { + CONF_LATITUDE: user_input.get( + CONF_LATITUDE, self.hass.config.latitude + ), + CONF_LONGITUDE: user_input.get( + CONF_LONGITUDE, self.hass.config.longitude + ), + } + ] + return self.async_create_entry( - title=f"Cloud API (API key: {user_input[CONF_API_KEY][:4]}...)", - data={ - CONF_API_KEY: user_input[CONF_API_KEY], - CONF_GEOGRAPHIES: [ - { - CONF_LATITUDE: user_input[CONF_LATITUDE], - CONF_LONGITUDE: user_input[CONF_LONGITUDE], - } - ], - }, + title=f"Cloud API (API key: {user_input[CONF_API_KEY][:4]}...)", data=data ) diff --git a/tests/components/airvisual/test_config_flow.py b/tests/components/airvisual/test_config_flow.py index fb4325bd6ee..5057f1c3345 100644 --- a/tests/components/airvisual/test_config_flow.py +++ b/tests/components/airvisual/test_config_flow.py @@ -1,6 +1,5 @@ """Define tests for the AirVisual config flow.""" -from unittest.mock import patch - +from asynctest import patch from pyairvisual.errors import InvalidKeyError from homeassistant import data_entry_flow @@ -8,7 +7,7 @@ from homeassistant.components.airvisual import CONF_GEOGRAPHIES, DOMAIN from homeassistant.config_entries import SOURCE_IMPORT, SOURCE_USER from homeassistant.const import CONF_API_KEY, CONF_LATITUDE, CONF_LONGITUDE -from tests.common import MockConfigEntry, mock_coro +from tests.common import MockConfigEntry async def test_duplicate_error(hass): @@ -30,8 +29,7 @@ async def test_invalid_api_key(hass): conf = {CONF_API_KEY: "abcde12345"} with patch( - "pyairvisual.api.API.nearest_city", - return_value=mock_coro(exception=InvalidKeyError), + "pyairvisual.api.API.nearest_city", side_effect=InvalidKeyError, ): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER}, data=conf @@ -53,16 +51,19 @@ async def test_step_import(hass): """Test that the import step works.""" conf = {CONF_API_KEY: "abcde12345"} - result = await hass.config_entries.flow.async_init( - DOMAIN, context={"source": SOURCE_IMPORT}, data=conf - ) + with patch( + "homeassistant.components.wwlln.async_setup_entry", return_value=True + ), patch("pyairvisual.api.API.nearest_city"): + result = await hass.config_entries.flow.async_init( + DOMAIN, context={"source": SOURCE_IMPORT}, data=conf + ) - assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY - assert result["title"] == "Cloud API (API key: abcd...)" - assert result["data"] == { - CONF_API_KEY: "abcde12345", - CONF_GEOGRAPHIES: [{CONF_LATITUDE: 32.87336, CONF_LONGITUDE: -117.22743}], - } + assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY + assert result["title"] == "Cloud API (API key: abcd...)" + assert result["data"] == { + CONF_API_KEY: "abcde12345", + CONF_GEOGRAPHIES: [{CONF_LATITUDE: 32.87336, CONF_LONGITUDE: -117.22743}], + } async def test_step_user(hass): @@ -74,8 +75,8 @@ async def test_step_user(hass): } with patch( - "pyairvisual.api.API.nearest_city", return_value=mock_coro(), - ): + "homeassistant.components.wwlln.async_setup_entry", return_value=True + ), patch("pyairvisual.api.API.nearest_city"): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER}, data=conf )