Store integration type in AirVisual config entry (#34621)

homeassistant/components/airvisual/__init__.py

@@ -30,6 +30,7 @@ from .const import (
     CONF_CITY,
     CONF_COUNTRY,
     CONF_GEOGRAPHIES,
+    CONF_INTEGRATION_TYPE,
     DATA_CLIENT,
     DOMAIN,
     INTEGRATION_TYPE_GEOGRAPHY,
@@ -120,7 +121,7 @@ async def async_setup(hass, config):
 
 @callback
 def _standardize_geography_config_entry(hass, config_entry):
-    """Ensure that geography observables have appropriate properties."""
+    """Ensure that geography config entries have appropriate properties."""
     entry_updates = {}
 
     if not config_entry.unique_id:
@@ -129,6 +130,30 @@ def _standardize_geography_config_entry(hass, config_entry):
     if not config_entry.options:
         # If the config entry doesn't already have any options set, set defaults:
         entry_updates["options"] = {CONF_SHOW_ON_MAP: True}
+    if CONF_INTEGRATION_TYPE not in config_entry.data:
+        # 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,
+        }
+
+    if not entry_updates:
+        return
+
+    hass.config_entries.async_update_entry(config_entry, **entry_updates)
+
+
+@callback
+def _standardize_node_pro_config_entry(hass, config_entry):
+    """Ensure that Node/Pro config entries have appropriate properties."""
+    entry_updates = {}
+
+    if CONF_INTEGRATION_TYPE not in config_entry.data:
+        # If the config entry data doesn't contain the integration type, add it:
+        entry_updates["data"] = {
+            **config_entry.data,
+            CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_NODE_PRO,
+        }
 
     if not entry_updates:
         return
@@ -151,6 +176,7 @@ async def async_setup_entry(hass, config_entry):
         # Only geography-based entries have options:
         config_entry.add_update_listener(async_update_options)
     else:
+        _standardize_node_pro_config_entry(hass, config_entry)
         airvisual = AirVisualNodeProData(hass, Client(websession), config_entry)
 
     await airvisual.async_update()

homeassistant/components/airvisual/config_flow.py

@@ -20,6 +20,7 @@ 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_INTEGRATION_TYPE,
     DOMAIN,
     INTEGRATION_TYPE_GEOGRAPHY,
     INTEGRATION_TYPE_NODE_PRO,
@@ -123,7 +124,8 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
                 checked_keys.add(user_input[CONF_API_KEY])
 
             return self.async_create_entry(
-                title=f"Cloud API ({geo_id})", data=user_input
+                title=f"Cloud API ({geo_id})",
+                data={**user_input, CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY},
             )
 
     async def async_step_import(self, import_config):
@@ -155,7 +157,8 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
             )
 
         return self.async_create_entry(
-            title=f"Node/Pro ({user_input[CONF_IP_ADDRESS]})", data=user_input
+            title=f"Node/Pro ({user_input[CONF_IP_ADDRESS]})",
+            data={**user_input, CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_NODE_PRO},
         )
 
     async def async_step_user(self, user_input=None):

homeassistant/components/airvisual/const.py

@@ -10,6 +10,7 @@ INTEGRATION_TYPE_NODE_PRO = "AirVisual Node/Pro"
 CONF_CITY = "city"
 CONF_COUNTRY = "country"
 CONF_GEOGRAPHIES = "geographies"
+CONF_INTEGRATION_TYPE = "integration_type"
 
 DATA_CLIENT = "client"
 

tests/components/airvisual/test_config_flow.py

@@ -5,6 +5,7 @@ from pyairvisual.errors import InvalidKeyError, NodeProError
 from homeassistant import data_entry_flow
 from homeassistant.components.airvisual import (
     CONF_GEOGRAPHIES,
+    CONF_INTEGRATION_TYPE,
     DOMAIN,
     INTEGRATION_TYPE_GEOGRAPHY,
     INTEGRATION_TYPE_NODE_PRO,
@@ -93,8 +94,8 @@ async def test_node_pro_error(hass):
         assert result["errors"] == {CONF_IP_ADDRESS: "unable_to_connect"}
 
 
-async def test_migration_1_2(hass):
+async def test_migration(hass):
-    """Test migrating from version 1 to version 2."""
+    """Test migrating from version 1 to the current version."""
     conf = {
         CONF_API_KEY: "abcde12345",
         CONF_GEOGRAPHIES: [
@@ -123,6 +124,7 @@ async def test_migration_1_2(hass):
         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"
@@ -131,6 +133,7 @@ async def test_migration_1_2(hass):
         CONF_API_KEY: "abcde12345",
         CONF_LATITUDE: 35.48847,
         CONF_LONGITUDE: 137.5263065,
+        CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY,
     }
 
 
@@ -186,6 +189,7 @@ async def test_step_geography(hass):
             CONF_API_KEY: "abcde12345",
             CONF_LATITUDE: 51.528308,
             CONF_LONGITUDE: -0.3817765,
+            CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY,
         }
 
 
@@ -207,6 +211,7 @@ async def test_step_node_pro(hass):
         assert result["data"] == {
             CONF_IP_ADDRESS: "192.168.1.100",
             CONF_PASSWORD: "my_password",
+            CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_NODE_PRO,
         }
 
 
@@ -231,6 +236,7 @@ async def test_step_import(hass):
             CONF_API_KEY: "abcde12345",
             CONF_LATITUDE: 51.528308,
             CONF_LONGITUDE: -0.3817765,
+            CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_GEOGRAPHY,
         }