Ensure AirVisual import config flow checks API key correctness (#32624)

* Add small improvements to AirVisual config flow tests

* Code review comments

* Code review comments

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])
+        return await self.async_step_user(import_config)
-
-        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,
-        )
 
     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]}...)",
+            title=f"Cloud API (API key: {user_input[CONF_API_KEY][:4]}...)", data=data
-            data={
-                CONF_API_KEY: user_input[CONF_API_KEY],
-                CONF_GEOGRAPHIES: [
-                    {
-                        CONF_LATITUDE: user_input[CONF_LATITUDE],
-                        CONF_LONGITUDE: user_input[CONF_LONGITUDE],
-                    }
-                ],
-            },
         )

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",
+        "pyairvisual.api.API.nearest_city", side_effect=InvalidKeyError,
-        return_value=mock_coro(exception=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(
+    with patch(
-        DOMAIN, context={"source": SOURCE_IMPORT}, data=conf
+        "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["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
-    assert result["title"] == "Cloud API (API key: abcd...)"
+        assert result["title"] == "Cloud API (API key: abcd...)"
-    assert result["data"] == {
+        assert result["data"] == {
-        CONF_API_KEY: "abcde12345",
+            CONF_API_KEY: "abcde12345",
-        CONF_GEOGRAPHIES: [{CONF_LATITUDE: 32.87336, CONF_LONGITUDE: -117.22743}],
+            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
         )