From 3bdf9628385a381d28f1ee092ab686eed23f8163 Mon Sep 17 00:00:00 2001 From: Aaron Bach Date: Mon, 1 Feb 2021 14:38:03 -0700 Subject: [PATCH] Add ability to configure AirVisual with city/state/country in UI (#44116) --- .../components/airvisual/__init__.py | 23 +- .../components/airvisual/config_flow.py | 162 +++++++----- homeassistant/components/airvisual/const.py | 3 +- homeassistant/components/airvisual/sensor.py | 45 +++- .../components/airvisual/strings.json | 22 +- .../components/airvisual/translations/en.json | 22 +- .../components/airvisual/test_config_flow.py | 237 +++++++++++++----- 7 files changed, 344 insertions(+), 170 deletions(-) diff --git a/homeassistant/components/airvisual/__init__.py b/homeassistant/components/airvisual/__init__.py index 956b168a665..3a88243b0b9 100644 --- a/homeassistant/components/airvisual/__init__.py +++ b/homeassistant/components/airvisual/__init__.py @@ -23,6 +23,7 @@ from homeassistant.const import ( CONF_STATE, ) from homeassistant.core import callback +from homeassistant.exceptions import ConfigEntryNotReady from homeassistant.helpers import aiohttp_client, config_validation as cv from homeassistant.helpers.update_coordinator import ( CoordinatorEntity, @@ -37,7 +38,7 @@ from .const import ( CONF_INTEGRATION_TYPE, DATA_COORDINATOR, DOMAIN, - INTEGRATION_TYPE_GEOGRAPHY, + INTEGRATION_TYPE_GEOGRAPHY_COORDS, INTEGRATION_TYPE_NODE_PRO, LOGGER, ) @@ -145,7 +146,7 @@ def _standardize_geography_config_entry(hass, config_entry): # If the config entry data doesn't contain the integration type, add it: entry_updates["data"] = { **config_entry.data, - CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY, + CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY_COORDS, } if not entry_updates: @@ -232,7 +233,6 @@ async def async_setup_entry(hass, config_entry): update_method=async_update_data, ) - hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id] = coordinator async_sync_geo_coordinator_update_intervals( hass, config_entry.data[CONF_API_KEY] ) @@ -262,9 +262,11 @@ async def async_setup_entry(hass, config_entry): update_method=async_update_data, ) - hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id] = coordinator - await coordinator.async_refresh() + if not coordinator.last_update_success: + raise ConfigEntryNotReady + + hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id] = coordinator for component in PLATFORMS: hass.async_create_task( @@ -299,10 +301,14 @@ async def async_migrate_entry(hass, config_entry): # For any geographies that remain, create a new config entry for each one: for geography in geographies: + if CONF_LATITUDE in geography: + source = "geography_by_coords" + else: + source = "geography_by_name" hass.async_create_task( hass.config_entries.flow.async_init( DOMAIN, - context={"source": "geography"}, + context={"source": source}, data={CONF_API_KEY: config_entry.data[CONF_API_KEY], **geography}, ) ) @@ -327,7 +333,10 @@ async def async_unload_entry(hass, config_entry): remove_listener = hass.data[DOMAIN][DATA_LISTENER].pop(config_entry.entry_id) remove_listener() - if config_entry.data[CONF_INTEGRATION_TYPE] == INTEGRATION_TYPE_GEOGRAPHY: + if ( + config_entry.data[CONF_INTEGRATION_TYPE] + == INTEGRATION_TYPE_GEOGRAPHY_COORDS + ): # Re-calculate the update interval period for any remaining consumers of # this API key: async_sync_geo_coordinator_update_intervals( diff --git a/homeassistant/components/airvisual/config_flow.py b/homeassistant/components/airvisual/config_flow.py index b086aeefc27..266f7b7c2c2 100644 --- a/homeassistant/components/airvisual/config_flow.py +++ b/homeassistant/components/airvisual/config_flow.py @@ -2,7 +2,12 @@ import asyncio from pyairvisual import CloudAPI, NodeSamba -from pyairvisual.errors import InvalidKeyError, NodeProError +from pyairvisual.errors import ( + AirVisualError, + InvalidKeyError, + NodeProError, + NotFoundError, +) import voluptuous as vol from homeassistant import config_entries @@ -13,20 +18,46 @@ from homeassistant.const import ( CONF_LONGITUDE, CONF_PASSWORD, CONF_SHOW_ON_MAP, + CONF_STATE, ) from homeassistant.core import callback from homeassistant.helpers import aiohttp_client, config_validation as cv from . import async_get_geography_id from .const import ( # pylint: disable=unused-import - CONF_GEOGRAPHIES, + CONF_CITY, + CONF_COUNTRY, CONF_INTEGRATION_TYPE, DOMAIN, - INTEGRATION_TYPE_GEOGRAPHY, + INTEGRATION_TYPE_GEOGRAPHY_COORDS, + INTEGRATION_TYPE_GEOGRAPHY_NAME, INTEGRATION_TYPE_NODE_PRO, LOGGER, ) +API_KEY_DATA_SCHEMA = vol.Schema({vol.Required(CONF_API_KEY): cv.string}) +GEOGRAPHY_NAME_SCHEMA = API_KEY_DATA_SCHEMA.extend( + { + vol.Required(CONF_CITY): cv.string, + vol.Required(CONF_STATE): cv.string, + vol.Required(CONF_COUNTRY): cv.string, + } +) +NODE_PRO_SCHEMA = vol.Schema( + {vol.Required(CONF_IP_ADDRESS): str, vol.Required(CONF_PASSWORD): cv.string} +) +PICK_INTEGRATION_TYPE_SCHEMA = vol.Schema( + { + vol.Required("type"): vol.In( + [ + INTEGRATION_TYPE_GEOGRAPHY_COORDS, + INTEGRATION_TYPE_GEOGRAPHY_NAME, + INTEGRATION_TYPE_NODE_PRO, + ] + ) + } +) + class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): """Handle an AirVisual config flow.""" @@ -36,16 +67,13 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): def __init__(self): """Initialize the config flow.""" + self._entry_data_for_reauth = None self._geo_id = None - self._latitude = None - self._longitude = None - - self.api_key_data_schema = vol.Schema({vol.Required(CONF_API_KEY): str}) @property - def geography_schema(self): + def geography_coords_schema(self): """Return the data schema for the cloud API.""" - return self.api_key_data_schema.extend( + return API_KEY_DATA_SCHEMA.extend( { vol.Required( CONF_LATITUDE, default=self.hass.config.latitude @@ -56,24 +84,6 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): } ) - @property - def pick_integration_type_schema(self): - """Return the data schema for picking the integration type.""" - return vol.Schema( - { - vol.Required("type"): vol.In( - [INTEGRATION_TYPE_GEOGRAPHY, INTEGRATION_TYPE_NODE_PRO] - ) - } - ) - - @property - def node_pro_schema(self): - """Return the data schema for a Node/Pro.""" - return vol.Schema( - {vol.Required(CONF_IP_ADDRESS): str, vol.Required(CONF_PASSWORD): str} - ) - async def _async_set_unique_id(self, unique_id): """Set the unique ID of the config flow and abort if it already exists.""" await self.async_set_unique_id(unique_id) @@ -85,33 +95,36 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): """Define the config flow to handle options.""" return AirVisualOptionsFlowHandler(config_entry) - async def async_step_geography(self, user_input=None): + async def async_step_geography(self, user_input, integration_type): """Handle the initialization of the integration via the cloud API.""" - if not user_input: - return self.async_show_form( - step_id="geography", data_schema=self.geography_schema - ) - self._geo_id = async_get_geography_id(user_input) await self._async_set_unique_id(self._geo_id) self._abort_if_unique_id_configured() + return await self.async_step_geography_finish(user_input, integration_type) - # Find older config entries without unique ID: - for entry in self._async_current_entries(): - if entry.version != 1: - continue + async def async_step_geography_by_coords(self, user_input=None): + """Handle the initialization of the cloud API based on latitude/longitude.""" + if not user_input: + return self.async_show_form( + step_id="geography_by_coords", data_schema=self.geography_coords_schema + ) - if any( - self._geo_id == async_get_geography_id(geography) - for geography in entry.data[CONF_GEOGRAPHIES] - ): - return self.async_abort(reason="already_configured") - - return await self.async_step_geography_finish( - user_input, "geography", self.geography_schema + return await self.async_step_geography( + user_input, INTEGRATION_TYPE_GEOGRAPHY_COORDS ) - async def async_step_geography_finish(self, user_input, error_step, error_schema): + async def async_step_geography_by_name(self, user_input=None): + """Handle the initialization of the cloud API based on city/state/country.""" + if not user_input: + return self.async_show_form( + step_id="geography_by_name", data_schema=GEOGRAPHY_NAME_SCHEMA + ) + + return await self.async_step_geography( + user_input, INTEGRATION_TYPE_GEOGRAPHY_NAME + ) + + async def async_step_geography_finish(self, user_input, integration_type): """Validate a Cloud API key.""" websession = aiohttp_client.async_get_clientsession(self.hass) cloud_api = CloudAPI(user_input[CONF_API_KEY], session=websession) @@ -123,16 +136,40 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): "airvisual_checked_api_keys_lock", asyncio.Lock() ) + if integration_type == INTEGRATION_TYPE_GEOGRAPHY_COORDS: + coro = cloud_api.air_quality.nearest_city() + error_schema = self.geography_coords_schema + error_step = "geography_by_coords" + else: + coro = cloud_api.air_quality.city( + user_input[CONF_CITY], user_input[CONF_STATE], user_input[CONF_COUNTRY] + ) + error_schema = GEOGRAPHY_NAME_SCHEMA + error_step = "geography_by_name" + async with valid_keys_lock: if user_input[CONF_API_KEY] not in valid_keys: try: - await cloud_api.air_quality.nearest_city() + await coro except InvalidKeyError: return self.async_show_form( step_id=error_step, data_schema=error_schema, errors={CONF_API_KEY: "invalid_api_key"}, ) + except NotFoundError: + return self.async_show_form( + step_id=error_step, + data_schema=error_schema, + errors={CONF_CITY: "location_not_found"}, + ) + except AirVisualError as err: + LOGGER.error(err) + return self.async_show_form( + step_id=error_step, + data_schema=error_schema, + errors={"base": "unknown"}, + ) valid_keys.add(user_input[CONF_API_KEY]) @@ -143,15 +180,13 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): return self.async_create_entry( title=f"Cloud API ({self._geo_id})", - data={**user_input, CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY}, + data={**user_input, CONF_INTEGRATION_TYPE: integration_type}, ) async def async_step_node_pro(self, user_input=None): """Handle the initialization of the integration with a Node/Pro.""" if not user_input: - return self.async_show_form( - step_id="node_pro", data_schema=self.node_pro_schema - ) + return self.async_show_form(step_id="node_pro", data_schema=NODE_PRO_SCHEMA) await self._async_set_unique_id(user_input[CONF_IP_ADDRESS]) @@ -163,7 +198,7 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): LOGGER.error("Error connecting to Node/Pro unit: %s", err) return self.async_show_form( step_id="node_pro", - data_schema=self.node_pro_schema, + data_schema=NODE_PRO_SCHEMA, errors={CONF_IP_ADDRESS: "cannot_connect"}, ) @@ -176,39 +211,34 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN): async def async_step_reauth(self, data): """Handle configuration by re-auth.""" + self._entry_data_for_reauth = data self._geo_id = async_get_geography_id(data) - self._latitude = data[CONF_LATITUDE] - self._longitude = data[CONF_LONGITUDE] - return await self.async_step_reauth_confirm() async def async_step_reauth_confirm(self, user_input=None): """Handle re-auth completion.""" if not user_input: return self.async_show_form( - step_id="reauth_confirm", data_schema=self.api_key_data_schema + step_id="reauth_confirm", data_schema=API_KEY_DATA_SCHEMA ) - conf = { - CONF_API_KEY: user_input[CONF_API_KEY], - CONF_LATITUDE: self._latitude, - CONF_LONGITUDE: self._longitude, - CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY, - } + conf = {CONF_API_KEY: user_input[CONF_API_KEY], **self._entry_data_for_reauth} return await self.async_step_geography_finish( - conf, "reauth_confirm", self.api_key_data_schema + conf, self._entry_data_for_reauth[CONF_INTEGRATION_TYPE] ) async def async_step_user(self, user_input=None): """Handle the start of the config flow.""" if not user_input: return self.async_show_form( - step_id="user", data_schema=self.pick_integration_type_schema + step_id="user", data_schema=PICK_INTEGRATION_TYPE_SCHEMA ) - if user_input["type"] == INTEGRATION_TYPE_GEOGRAPHY: - return await self.async_step_geography() + if user_input["type"] == INTEGRATION_TYPE_GEOGRAPHY_COORDS: + return await self.async_step_geography_by_coords() + if user_input["type"] == INTEGRATION_TYPE_GEOGRAPHY_NAME: + return await self.async_step_geography_by_name() return await self.async_step_node_pro() diff --git a/homeassistant/components/airvisual/const.py b/homeassistant/components/airvisual/const.py index a98a899b762..510ada2b68c 100644 --- a/homeassistant/components/airvisual/const.py +++ b/homeassistant/components/airvisual/const.py @@ -4,7 +4,8 @@ import logging DOMAIN = "airvisual" LOGGER = logging.getLogger(__package__) -INTEGRATION_TYPE_GEOGRAPHY = "Geographical Location" +INTEGRATION_TYPE_GEOGRAPHY_COORDS = "Geographical Location by Latitude/Longitude" +INTEGRATION_TYPE_GEOGRAPHY_NAME = "Geographical Location by Name" INTEGRATION_TYPE_NODE_PRO = "AirVisual Node/Pro" CONF_CITY = "city" diff --git a/homeassistant/components/airvisual/sensor.py b/homeassistant/components/airvisual/sensor.py index ae9995f36c3..680059af411 100644 --- a/homeassistant/components/airvisual/sensor.py +++ b/homeassistant/components/airvisual/sensor.py @@ -27,7 +27,8 @@ from .const import ( CONF_INTEGRATION_TYPE, DATA_COORDINATOR, DOMAIN, - INTEGRATION_TYPE_GEOGRAPHY, + INTEGRATION_TYPE_GEOGRAPHY_COORDS, + INTEGRATION_TYPE_GEOGRAPHY_NAME, ) _LOGGER = getLogger(__name__) @@ -115,7 +116,10 @@ async def async_setup_entry(hass, config_entry, async_add_entities): """Set up AirVisual sensors based on a config entry.""" coordinator = hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id] - if config_entry.data[CONF_INTEGRATION_TYPE] == INTEGRATION_TYPE_GEOGRAPHY: + if config_entry.data[CONF_INTEGRATION_TYPE] in [ + INTEGRATION_TYPE_GEOGRAPHY_COORDS, + INTEGRATION_TYPE_GEOGRAPHY_NAME, + ]: sensors = [ AirVisualGeographySensor( coordinator, @@ -208,17 +212,32 @@ class AirVisualGeographySensor(AirVisualEntity): } ) - if CONF_LATITUDE in self._config_entry.data: - if self._config_entry.options[CONF_SHOW_ON_MAP]: - self._attrs[ATTR_LATITUDE] = self._config_entry.data[CONF_LATITUDE] - self._attrs[ATTR_LONGITUDE] = self._config_entry.data[CONF_LONGITUDE] - self._attrs.pop("lati", None) - self._attrs.pop("long", None) - else: - self._attrs["lati"] = self._config_entry.data[CONF_LATITUDE] - self._attrs["long"] = self._config_entry.data[CONF_LONGITUDE] - self._attrs.pop(ATTR_LATITUDE, None) - self._attrs.pop(ATTR_LONGITUDE, None) + # Displaying the geography on the map relies upon putting the latitude/longitude + # in the entity attributes with "latitude" and "longitude" as the keys. + # Conversely, we can hide the location on the map by using other keys, like + # "lati" and "long". + # + # We use any coordinates in the config entry and, in the case of a geography by + # name, we fall back to the latitude longitude provided in the coordinator data: + latitude = self._config_entry.data.get( + CONF_LATITUDE, + self.coordinator.data["location"]["coordinates"][1], + ) + longitude = self._config_entry.data.get( + CONF_LONGITUDE, + self.coordinator.data["location"]["coordinates"][0], + ) + + if self._config_entry.options[CONF_SHOW_ON_MAP]: + self._attrs[ATTR_LATITUDE] = latitude + self._attrs[ATTR_LONGITUDE] = longitude + self._attrs.pop("lati", None) + self._attrs.pop("long", None) + else: + self._attrs["lati"] = latitude + self._attrs["long"] = longitude + self._attrs.pop(ATTR_LATITUDE, None) + self._attrs.pop(ATTR_LONGITUDE, None) class AirVisualNodeProSensor(AirVisualEntity): diff --git a/homeassistant/components/airvisual/strings.json b/homeassistant/components/airvisual/strings.json index 22f9c80f313..8d2dce85a17 100644 --- a/homeassistant/components/airvisual/strings.json +++ b/homeassistant/components/airvisual/strings.json @@ -1,15 +1,25 @@ { "config": { "step": { - "geography": { + "geography_by_coords": { "title": "Configure a Geography", - "description": "Use the AirVisual cloud API to monitor a geographical location.", + "description": "Use the AirVisual cloud API to monitor a latitude/longitude.", "data": { "api_key": "[%key:common::config_flow::data::api_key%]", "latitude": "[%key:common::config_flow::data::latitude%]", "longitude": "[%key:common::config_flow::data::longitude%]" } }, + "geography_by_name": { + "title": "Configure a Geography", + "description": "Use the AirVisual cloud API to monitor a city/state/country.", + "data": { + "api_key": "[%key:common::config_flow::data::api_key%]", + "city": "City", + "country": "Country", + "state": "state" + } + }, "node_pro": { "title": "Configure an AirVisual Node/Pro", "description": "Monitor a personal AirVisual unit. The password can be retrieved from the unit's UI.", @@ -26,17 +36,13 @@ }, "user": { "title": "Configure AirVisual", - "description": "Pick what type of AirVisual data you want to monitor.", - "data": { - "cloud_api": "Geographical Location", - "node_pro": "AirVisual Node Pro", - "type": "Integration Type" - } + "description": "Pick what type of AirVisual data you want to monitor." } }, "error": { "general_error": "[%key:common::config_flow::error::unknown%]", "invalid_api_key": "[%key:common::config_flow::error::invalid_api_key%]", + "location_not_found": "Location not found", "cannot_connect": "[%key:common::config_flow::error::cannot_connect%]" }, "abort": { diff --git a/homeassistant/components/airvisual/translations/en.json b/homeassistant/components/airvisual/translations/en.json index 129abcc29e5..1e3cb59a520 100644 --- a/homeassistant/components/airvisual/translations/en.json +++ b/homeassistant/components/airvisual/translations/en.json @@ -7,16 +7,27 @@ "error": { "cannot_connect": "Failed to connect", "general_error": "Unexpected error", - "invalid_api_key": "Invalid API key" + "invalid_api_key": "Invalid API key", + "location_not_found": "Location not found" }, "step": { - "geography": { + "geography_by_coords": { "data": { "api_key": "API Key", "latitude": "Latitude", "longitude": "Longitude" }, - "description": "Use the AirVisual cloud API to monitor a geographical location.", + "description": "Use the AirVisual cloud API to monitor a latitude/longitude.", + "title": "Configure a Geography" + }, + "geography_by_name": { + "data": { + "api_key": "API Key", + "city": "City", + "country": "Country", + "state": "state" + }, + "description": "Use the AirVisual cloud API to monitor a city/state/country.", "title": "Configure a Geography" }, "node_pro": { @@ -34,11 +45,6 @@ "title": "Re-authenticate AirVisual" }, "user": { - "data": { - "cloud_api": "Geographical Location", - "node_pro": "AirVisual Node Pro", - "type": "Integration Type" - }, "description": "Pick what type of AirVisual data you want to monitor.", "title": "Configure AirVisual" } diff --git a/tests/components/airvisual/test_config_flow.py b/tests/components/airvisual/test_config_flow.py index 4e550d94b09..248abaf6b5f 100644 --- a/tests/components/airvisual/test_config_flow.py +++ b/tests/components/airvisual/test_config_flow.py @@ -1,14 +1,22 @@ """Define tests for the AirVisual config flow.""" from unittest.mock import patch -from pyairvisual.errors import InvalidKeyError, NodeProError +from pyairvisual.errors import ( + AirVisualError, + InvalidKeyError, + NodeProError, + NotFoundError, +) from homeassistant import data_entry_flow -from homeassistant.components.airvisual import ( +from homeassistant.components.airvisual.const import ( + CONF_CITY, + CONF_COUNTRY, CONF_GEOGRAPHIES, CONF_INTEGRATION_TYPE, DOMAIN, - INTEGRATION_TYPE_GEOGRAPHY, + INTEGRATION_TYPE_GEOGRAPHY_COORDS, + INTEGRATION_TYPE_GEOGRAPHY_NAME, INTEGRATION_TYPE_NODE_PRO, ) from homeassistant.config_entries import SOURCE_USER @@ -19,6 +27,7 @@ from homeassistant.const import ( CONF_LONGITUDE, CONF_PASSWORD, CONF_SHOW_ON_MAP, + CONF_STATE, ) from homeassistant.setup import async_setup_component @@ -38,7 +47,9 @@ async def test_duplicate_error(hass): ).add_to_hass(hass) result = await hass.config_entries.flow.async_init( - DOMAIN, context={"source": SOURCE_USER}, data={"type": "Geographical Location"} + DOMAIN, + context={"source": SOURCE_USER}, + data={"type": INTEGRATION_TYPE_GEOGRAPHY_COORDS}, ) result = await hass.config_entries.flow.async_configure( result["flow_id"], user_input=geography_conf @@ -64,14 +75,8 @@ async def test_duplicate_error(hass): assert result["reason"] == "already_configured" -async def test_invalid_identifier(hass): - """Test that an invalid API key or Node/Pro ID throws an error.""" - geography_conf = { - CONF_API_KEY: "abcde12345", - CONF_LATITUDE: 51.528308, - CONF_LONGITUDE: -0.3817765, - } - +async def test_invalid_identifier_geography_api_key(hass): + """Test that an invalid API key throws an error.""" with patch( "pyairvisual.air_quality.AirQuality.nearest_city", side_effect=InvalidKeyError, @@ -79,64 +84,73 @@ async def test_invalid_identifier(hass): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER}, - data={"type": "Geographical Location"}, + data={"type": INTEGRATION_TYPE_GEOGRAPHY_COORDS}, ) result = await hass.config_entries.flow.async_configure( - result["flow_id"], user_input=geography_conf + result["flow_id"], + user_input={ + CONF_API_KEY: "abcde12345", + CONF_LATITUDE: 51.528308, + CONF_LONGITUDE: -0.3817765, + }, ) assert result["type"] == data_entry_flow.RESULT_TYPE_FORM assert result["errors"] == {CONF_API_KEY: "invalid_api_key"} -async def test_migration(hass): - """Test migrating from version 1 to the current version.""" - conf = { - CONF_API_KEY: "abcde12345", - CONF_GEOGRAPHIES: [ - {CONF_LATITUDE: 51.528308, CONF_LONGITUDE: -0.3817765}, - {CONF_LATITUDE: 35.48847, CONF_LONGITUDE: 137.5263065}, - ], - } - - config_entry = MockConfigEntry( - domain=DOMAIN, version=1, unique_id="abcde12345", data=conf - ) - config_entry.add_to_hass(hass) - - assert len(hass.config_entries.async_entries(DOMAIN)) == 1 - - with patch("pyairvisual.air_quality.AirQuality.nearest_city"), patch.object( - hass.config_entries, "async_forward_entry_setup" +async def test_invalid_identifier_geography_name(hass): + """Test that an invalid location name throws an error.""" + with patch( + "pyairvisual.air_quality.AirQuality.city", + side_effect=NotFoundError, ): - assert await async_setup_component(hass, DOMAIN, {DOMAIN: conf}) - await hass.async_block_till_done() + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": SOURCE_USER}, + data={"type": INTEGRATION_TYPE_GEOGRAPHY_NAME}, + ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ + CONF_API_KEY: "abcde12345", + CONF_CITY: "Beijing", + CONF_STATE: "Beijing", + CONF_COUNTRY: "China", + }, + ) - config_entries = hass.config_entries.async_entries(DOMAIN) - - assert len(config_entries) == 2 - - assert config_entries[0].unique_id == "51.528308, -0.3817765" - assert config_entries[0].title == "Cloud API (51.528308, -0.3817765)" - assert config_entries[0].data == { - CONF_API_KEY: "abcde12345", - CONF_LATITUDE: 51.528308, - CONF_LONGITUDE: -0.3817765, - CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY, - } - - assert config_entries[1].unique_id == "35.48847, 137.5263065" - assert config_entries[1].title == "Cloud API (35.48847, 137.5263065)" - assert config_entries[1].data == { - CONF_API_KEY: "abcde12345", - CONF_LATITUDE: 35.48847, - CONF_LONGITUDE: 137.5263065, - CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY, - } + assert result["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["errors"] == {CONF_CITY: "location_not_found"} -async def test_node_pro_error(hass): - """Test that an invalid Node/Pro ID shows an error.""" +async def test_invalid_identifier_geography_unknown(hass): + """Test that an unknown identifier issue throws an error.""" + with patch( + "pyairvisual.air_quality.AirQuality.city", + side_effect=AirVisualError, + ): + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": SOURCE_USER}, + data={"type": INTEGRATION_TYPE_GEOGRAPHY_NAME}, + ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], + user_input={ + CONF_API_KEY: "abcde12345", + CONF_CITY: "Beijing", + CONF_STATE: "Beijing", + CONF_COUNTRY: "China", + }, + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["errors"] == {"base": "unknown"} + + +async def test_invalid_identifier_node_pro(hass): + """Test that an invalid Node/Pro identifier shows an error.""" node_pro_conf = {CONF_IP_ADDRESS: "192.168.1.100", CONF_PASSWORD: "my_password"} with patch( @@ -153,6 +167,53 @@ async def test_node_pro_error(hass): assert result["errors"] == {CONF_IP_ADDRESS: "cannot_connect"} +async def test_migration(hass): + """Test migrating from version 1 to the current version.""" + conf = { + CONF_API_KEY: "abcde12345", + CONF_GEOGRAPHIES: [ + {CONF_LATITUDE: 51.528308, CONF_LONGITUDE: -0.3817765}, + {CONF_CITY: "Beijing", CONF_STATE: "Beijing", CONF_COUNTRY: "China"}, + ], + } + + config_entry = MockConfigEntry( + domain=DOMAIN, version=1, unique_id="abcde12345", data=conf + ) + config_entry.add_to_hass(hass) + + assert len(hass.config_entries.async_entries(DOMAIN)) == 1 + + with patch("pyairvisual.air_quality.AirQuality.city"), patch( + "pyairvisual.air_quality.AirQuality.nearest_city" + ), patch.object(hass.config_entries, "async_forward_entry_setup"): + assert await async_setup_component(hass, DOMAIN, {DOMAIN: conf}) + await hass.async_block_till_done() + + config_entries = hass.config_entries.async_entries(DOMAIN) + + assert len(config_entries) == 2 + + assert config_entries[0].unique_id == "51.528308, -0.3817765" + assert config_entries[0].title == "Cloud API (51.528308, -0.3817765)" + assert config_entries[0].data == { + CONF_API_KEY: "abcde12345", + CONF_LATITUDE: 51.528308, + CONF_LONGITUDE: -0.3817765, + CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY_COORDS, + } + + assert config_entries[1].unique_id == "Beijing, Beijing, China" + assert config_entries[1].title == "Cloud API (Beijing, Beijing, China)" + assert config_entries[1].data == { + CONF_API_KEY: "abcde12345", + CONF_CITY: "Beijing", + CONF_STATE: "Beijing", + CONF_COUNTRY: "China", + CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY_NAME, + } + + async def test_options_flow(hass): """Test config flow options.""" geography_conf = { @@ -186,8 +247,8 @@ async def test_options_flow(hass): assert config_entry.options == {CONF_SHOW_ON_MAP: False} -async def test_step_geography(hass): - """Test the geograph (cloud API) step.""" +async def test_step_geography_by_coords(hass): + """Test setting up a geopgraphy entry by latitude/longitude.""" conf = { CONF_API_KEY: "abcde12345", CONF_LATITUDE: 51.528308, @@ -200,7 +261,7 @@ async def test_step_geography(hass): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER}, - data={"type": "Geographical Location"}, + data={"type": INTEGRATION_TYPE_GEOGRAPHY_COORDS}, ) result = await hass.config_entries.flow.async_configure( result["flow_id"], user_input=conf @@ -212,7 +273,39 @@ async def test_step_geography(hass): CONF_API_KEY: "abcde12345", CONF_LATITUDE: 51.528308, CONF_LONGITUDE: -0.3817765, - CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY, + CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY_COORDS, + } + + +async def test_step_geography_by_name(hass): + """Test setting up a geopgraphy entry by city/state/country.""" + conf = { + CONF_API_KEY: "abcde12345", + CONF_CITY: "Beijing", + CONF_STATE: "Beijing", + CONF_COUNTRY: "China", + } + + with patch( + "homeassistant.components.airvisual.async_setup_entry", return_value=True + ), patch("pyairvisual.air_quality.AirQuality.city"): + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": SOURCE_USER}, + data={"type": INTEGRATION_TYPE_GEOGRAPHY_NAME}, + ) + result = await hass.config_entries.flow.async_configure( + result["flow_id"], user_input=conf + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY + assert result["title"] == "Cloud API (Beijing, Beijing, China)" + assert result["data"] == { + CONF_API_KEY: "abcde12345", + CONF_CITY: "Beijing", + CONF_STATE: "Beijing", + CONF_COUNTRY: "China", + CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY_NAME, } @@ -244,18 +337,19 @@ async def test_step_node_pro(hass): async def test_step_reauth(hass): """Test that the reauth step works.""" - geography_conf = { + entry_data = { CONF_API_KEY: "abcde12345", CONF_LATITUDE: 51.528308, CONF_LONGITUDE: -0.3817765, + CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY_COORDS, } MockConfigEntry( - domain=DOMAIN, unique_id="51.528308, -0.3817765", data=geography_conf + domain=DOMAIN, unique_id="51.528308, -0.3817765", data=entry_data ).add_to_hass(hass) result = await hass.config_entries.flow.async_init( - DOMAIN, context={"source": "reauth"}, data=geography_conf + DOMAIN, context={"source": "reauth"}, data=entry_data ) assert result["step_id"] == "reauth_confirm" @@ -287,11 +381,20 @@ async def test_step_user(hass): result = await hass.config_entries.flow.async_init( DOMAIN, context={"source": SOURCE_USER}, - data={"type": INTEGRATION_TYPE_GEOGRAPHY}, + data={"type": INTEGRATION_TYPE_GEOGRAPHY_COORDS}, ) assert result["type"] == data_entry_flow.RESULT_TYPE_FORM - assert result["step_id"] == "geography" + assert result["step_id"] == "geography_by_coords" + + result = await hass.config_entries.flow.async_init( + DOMAIN, + context={"source": SOURCE_USER}, + data={"type": INTEGRATION_TYPE_GEOGRAPHY_NAME}, + ) + + assert result["type"] == data_entry_flow.RESULT_TYPE_FORM + assert result["step_id"] == "geography_by_name" result = await hass.config_entries.flow.async_init( DOMAIN,