Use current config entry standards for AirVisual (#57132)

This commit is contained in:
Aaron Bach 2021-10-11 09:17:43 -06:00 committed by GitHub
parent 6c470ac28b
commit b72f1553ea
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 81 additions and 89 deletions

View File

@ -106,8 +106,8 @@ def async_get_cloud_coordinators_by_api_key(
"""Get all DataUpdateCoordinator objects related to a particular API key.""" """Get all DataUpdateCoordinator objects related to a particular API key."""
coordinators = [] coordinators = []
for entry_id, coordinator in hass.data[DOMAIN][DATA_COORDINATOR].items(): for entry_id, coordinator in hass.data[DOMAIN][DATA_COORDINATOR].items():
config_entry = hass.config_entries.async_get_entry(entry_id) entry = hass.config_entries.async_get_entry(entry_id)
if config_entry and config_entry.data.get(CONF_API_KEY) == api_key: if entry and entry.data.get(CONF_API_KEY) == api_key:
coordinators.append(coordinator) coordinators.append(coordinator)
return coordinators return coordinators
@ -137,25 +137,25 @@ def async_sync_geo_coordinator_update_intervals(
@callback @callback
def _standardize_geography_config_entry( def _standardize_geography_config_entry(
hass: HomeAssistant, config_entry: ConfigEntry hass: HomeAssistant, entry: ConfigEntry
) -> None: ) -> None:
"""Ensure that geography config entries have appropriate properties.""" """Ensure that geography config entries have appropriate properties."""
entry_updates = {} entry_updates = {}
if not config_entry.unique_id: if not entry.unique_id:
# If the config entry doesn't already have a unique ID, set one: # If the config entry doesn't already have a unique ID, set one:
entry_updates["unique_id"] = config_entry.data[CONF_API_KEY] entry_updates["unique_id"] = entry.data[CONF_API_KEY]
if not config_entry.options: if not entry.options:
# If the config entry doesn't already have any options set, set defaults: # If the config entry doesn't already have any options set, set defaults:
entry_updates["options"] = {CONF_SHOW_ON_MAP: True} entry_updates["options"] = {CONF_SHOW_ON_MAP: True}
if config_entry.data.get(CONF_INTEGRATION_TYPE) not in [ if entry.data.get(CONF_INTEGRATION_TYPE) not in [
INTEGRATION_TYPE_GEOGRAPHY_COORDS, INTEGRATION_TYPE_GEOGRAPHY_COORDS,
INTEGRATION_TYPE_GEOGRAPHY_NAME, INTEGRATION_TYPE_GEOGRAPHY_NAME,
]: ]:
# If the config entry data doesn't contain an integration type that we know # If the config entry data doesn't contain an integration type that we know
# about, infer it from the data we have: # about, infer it from the data we have:
entry_updates["data"] = {**config_entry.data} entry_updates["data"] = {**entry.data}
if CONF_CITY in config_entry.data: if CONF_CITY in entry.data:
entry_updates["data"][ entry_updates["data"][
CONF_INTEGRATION_TYPE CONF_INTEGRATION_TYPE
] = INTEGRATION_TYPE_GEOGRAPHY_NAME ] = INTEGRATION_TYPE_GEOGRAPHY_NAME
@ -167,51 +167,49 @@ def _standardize_geography_config_entry(
if not entry_updates: if not entry_updates:
return return
hass.config_entries.async_update_entry(config_entry, **entry_updates) hass.config_entries.async_update_entry(entry, **entry_updates)
@callback @callback
def _standardize_node_pro_config_entry( def _standardize_node_pro_config_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
hass: HomeAssistant, config_entry: ConfigEntry
) -> None:
"""Ensure that Node/Pro config entries have appropriate properties.""" """Ensure that Node/Pro config entries have appropriate properties."""
entry_updates: dict[str, Any] = {} entry_updates: dict[str, Any] = {}
if CONF_INTEGRATION_TYPE not in config_entry.data: if CONF_INTEGRATION_TYPE not in entry.data:
# If the config entry data doesn't contain the integration type, add it: # If the config entry data doesn't contain the integration type, add it:
entry_updates["data"] = { entry_updates["data"] = {
**config_entry.data, **entry.data,
CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_NODE_PRO, CONF_INTEGRATION_TYPE: INTEGRATION_TYPE_NODE_PRO,
} }
if not entry_updates: if not entry_updates:
return return
hass.config_entries.async_update_entry(config_entry, **entry_updates) hass.config_entries.async_update_entry(entry, **entry_updates)
async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Set up AirVisual as config entry.""" """Set up AirVisual as config entry."""
hass.data.setdefault(DOMAIN, {DATA_COORDINATOR: {}}) hass.data.setdefault(DOMAIN, {DATA_COORDINATOR: {}})
if CONF_API_KEY in config_entry.data: if CONF_API_KEY in entry.data:
_standardize_geography_config_entry(hass, config_entry) _standardize_geography_config_entry(hass, entry)
websession = aiohttp_client.async_get_clientsession(hass) websession = aiohttp_client.async_get_clientsession(hass)
cloud_api = CloudAPI(config_entry.data[CONF_API_KEY], session=websession) cloud_api = CloudAPI(entry.data[CONF_API_KEY], session=websession)
async def async_update_data() -> dict[str, Any]: async def async_update_data() -> dict[str, Any]:
"""Get new data from the API.""" """Get new data from the API."""
if CONF_CITY in config_entry.data: if CONF_CITY in entry.data:
api_coro = cloud_api.air_quality.city( api_coro = cloud_api.air_quality.city(
config_entry.data[CONF_CITY], entry.data[CONF_CITY],
config_entry.data[CONF_STATE], entry.data[CONF_STATE],
config_entry.data[CONF_COUNTRY], entry.data[CONF_COUNTRY],
) )
else: else:
api_coro = cloud_api.air_quality.nearest_city( api_coro = cloud_api.air_quality.nearest_city(
config_entry.data[CONF_LATITUDE], entry.data[CONF_LATITUDE],
config_entry.data[CONF_LONGITUDE], entry.data[CONF_LONGITUDE],
) )
try: try:
@ -225,7 +223,7 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
coordinator = DataUpdateCoordinator( coordinator = DataUpdateCoordinator(
hass, hass,
LOGGER, LOGGER,
name=async_get_geography_id(config_entry.data), name=async_get_geography_id(entry.data),
# We give a placeholder update interval in order to create the coordinator; # We give a placeholder update interval in order to create the coordinator;
# then, below, we use the coordinator's presence (along with any other # then, below, we use the coordinator's presence (along with any other
# coordinators using the same API key) to calculate an actual, leveled # coordinators using the same API key) to calculate an actual, leveled
@ -235,16 +233,14 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
) )
# Only geography-based entries have options: # Only geography-based entries have options:
config_entry.async_on_unload( entry.async_on_unload(entry.add_update_listener(async_reload_entry))
config_entry.add_update_listener(async_reload_entry)
)
else: else:
# Remove outdated air_quality entities from the entity registry if they exist: # Remove outdated air_quality entities from the entity registry if they exist:
ent_reg = entity_registry.async_get(hass) ent_reg = entity_registry.async_get(hass)
for entity_entry in [ for entity_entry in [
e e
for e in ent_reg.entities.values() for e in ent_reg.entities.values()
if e.config_entry_id == config_entry.entry_id if e.config_entry_id == entry.entry_id
and e.entity_id.startswith("air_quality") and e.entity_id.startswith("air_quality")
]: ]:
LOGGER.debug( LOGGER.debug(
@ -252,13 +248,13 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
) )
ent_reg.async_remove(entity_entry.entity_id) ent_reg.async_remove(entity_entry.entity_id)
_standardize_node_pro_config_entry(hass, config_entry) _standardize_node_pro_config_entry(hass, entry)
async def async_update_data() -> dict[str, Any]: async def async_update_data() -> dict[str, Any]:
"""Get new data from the API.""" """Get new data from the API."""
try: try:
async with NodeSamba( async with NodeSamba(
config_entry.data[CONF_IP_ADDRESS], config_entry.data[CONF_PASSWORD] entry.data[CONF_IP_ADDRESS], entry.data[CONF_PASSWORD]
) as node: ) as node:
data = await node.async_get_latest_measurements() data = await node.async_get_latest_measurements()
return cast(Dict[str, Any], data) return cast(Dict[str, Any], data)
@ -275,40 +271,38 @@ async def async_setup_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> b
await coordinator.async_config_entry_first_refresh() await coordinator.async_config_entry_first_refresh()
hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id] = coordinator hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id] = coordinator
# Reassess the interval between 2 server requests # Reassess the interval between 2 server requests
if CONF_API_KEY in config_entry.data: if CONF_API_KEY in entry.data:
async_sync_geo_coordinator_update_intervals( async_sync_geo_coordinator_update_intervals(hass, entry.data[CONF_API_KEY])
hass, config_entry.data[CONF_API_KEY]
)
hass.config_entries.async_setup_platforms(config_entry, PLATFORMS) hass.config_entries.async_setup_platforms(entry, PLATFORMS)
return True return True
async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: async def async_migrate_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Migrate an old config entry.""" """Migrate an old config entry."""
version = config_entry.version version = entry.version
LOGGER.debug("Migrating from version %s", version) LOGGER.debug("Migrating from version %s", version)
# 1 -> 2: One geography per config entry # 1 -> 2: One geography per config entry
if version == 1: if version == 1:
version = config_entry.version = 2 version = entry.version = 2
# Update the config entry to only include the first geography (there is always # Update the config entry to only include the first geography (there is always
# guaranteed to be at least one): # guaranteed to be at least one):
geographies = list(config_entry.data[CONF_GEOGRAPHIES]) geographies = list(entry.data[CONF_GEOGRAPHIES])
first_geography = geographies.pop(0) first_geography = geographies.pop(0)
first_id = async_get_geography_id(first_geography) first_id = async_get_geography_id(first_geography)
hass.config_entries.async_update_entry( hass.config_entries.async_update_entry(
config_entry, entry,
unique_id=first_id, unique_id=first_id,
title=f"Cloud API ({first_id})", title=f"Cloud API ({first_id})",
data={CONF_API_KEY: config_entry.data[CONF_API_KEY], **first_geography}, data={CONF_API_KEY: entry.data[CONF_API_KEY], **first_geography},
) )
# For any geographies that remain, create a new config entry for each one: # For any geographies that remain, create a new config entry for each one:
@ -321,7 +315,7 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
hass.config_entries.flow.async_init( hass.config_entries.flow.async_init(
DOMAIN, DOMAIN,
context={"source": source}, context={"source": source},
data={CONF_API_KEY: config_entry.data[CONF_API_KEY], **geography}, data={CONF_API_KEY: entry.data[CONF_API_KEY], **geography},
) )
) )
@ -330,40 +324,40 @@ async def async_migrate_entry(hass: HomeAssistant, config_entry: ConfigEntry) ->
return True return True
async def async_unload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> bool: async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
"""Unload an AirVisual config entry.""" """Unload an AirVisual config entry."""
unload_ok = await hass.config_entries.async_unload_platforms( unload_ok = await hass.config_entries.async_unload_platforms(entry, PLATFORMS)
config_entry, PLATFORMS
)
if unload_ok: if unload_ok:
hass.data[DOMAIN][DATA_COORDINATOR].pop(config_entry.entry_id) hass.data[DOMAIN][DATA_COORDINATOR].pop(entry.entry_id)
if CONF_API_KEY in config_entry.data: if CONF_API_KEY in entry.data:
# Re-calculate the update interval period for any remaining consumers of # Re-calculate the update interval period for any remaining consumers of
# this API key: # this API key:
async_sync_geo_coordinator_update_intervals( async_sync_geo_coordinator_update_intervals(hass, entry.data[CONF_API_KEY])
hass, config_entry.data[CONF_API_KEY]
)
return unload_ok return unload_ok
async def async_reload_entry(hass: HomeAssistant, config_entry: ConfigEntry) -> None: async def async_reload_entry(hass: HomeAssistant, entry: ConfigEntry) -> None:
"""Handle an options update.""" """Handle an options update."""
await hass.config_entries.async_reload(config_entry.entry_id) await hass.config_entries.async_reload(entry.entry_id)
class AirVisualEntity(CoordinatorEntity): class AirVisualEntity(CoordinatorEntity):
"""Define a generic AirVisual entity.""" """Define a generic AirVisual entity."""
def __init__( def __init__(
self, coordinator: DataUpdateCoordinator, description: EntityDescription self,
coordinator: DataUpdateCoordinator,
entry: ConfigEntry,
description: EntityDescription,
) -> None: ) -> None:
"""Initialize.""" """Initialize."""
super().__init__(coordinator) super().__init__(coordinator)
self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: DEFAULT_ATTRIBUTION} self._attr_extra_state_attributes = {ATTR_ATTRIBUTION: DEFAULT_ATTRIBUTION}
self._entry = entry
self.entity_description = description self.entity_description = description
async def async_added_to_hass(self) -> None: async def async_added_to_hass(self) -> None:

View File

@ -262,9 +262,9 @@ class AirVisualFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
class AirVisualOptionsFlowHandler(config_entries.OptionsFlow): class AirVisualOptionsFlowHandler(config_entries.OptionsFlow):
"""Handle an AirVisual options flow.""" """Handle an AirVisual options flow."""
def __init__(self, config_entry: ConfigEntry) -> None: def __init__(self, entry: ConfigEntry) -> None:
"""Initialize.""" """Initialize."""
self.config_entry = config_entry self.entry = entry
async def async_step_init( async def async_step_init(
self, user_input: dict[str, str] | None = None self, user_input: dict[str, str] | None = None
@ -279,7 +279,7 @@ class AirVisualOptionsFlowHandler(config_entries.OptionsFlow):
{ {
vol.Required( vol.Required(
CONF_SHOW_ON_MAP, CONF_SHOW_ON_MAP,
default=self.config_entry.options.get(CONF_SHOW_ON_MAP), default=self.entry.options.get(CONF_SHOW_ON_MAP),
): bool ): bool
} }
), ),

View File

@ -189,26 +189,24 @@ POLLUTANT_UNITS = {
async def async_setup_entry( async def async_setup_entry(
hass: HomeAssistant, hass: HomeAssistant, entry: ConfigEntry, async_add_entities: AddEntitiesCallback
config_entry: ConfigEntry,
async_add_entities: AddEntitiesCallback,
) -> None: ) -> None:
"""Set up AirVisual sensors based on a config entry.""" """Set up AirVisual sensors based on a config entry."""
coordinator = hass.data[DOMAIN][DATA_COORDINATOR][config_entry.entry_id] coordinator = hass.data[DOMAIN][DATA_COORDINATOR][entry.entry_id]
sensors: list[AirVisualGeographySensor | AirVisualNodeProSensor] sensors: list[AirVisualGeographySensor | AirVisualNodeProSensor]
if config_entry.data[CONF_INTEGRATION_TYPE] in ( if entry.data[CONF_INTEGRATION_TYPE] in (
INTEGRATION_TYPE_GEOGRAPHY_COORDS, INTEGRATION_TYPE_GEOGRAPHY_COORDS,
INTEGRATION_TYPE_GEOGRAPHY_NAME, INTEGRATION_TYPE_GEOGRAPHY_NAME,
): ):
sensors = [ sensors = [
AirVisualGeographySensor(coordinator, config_entry, description, locale) AirVisualGeographySensor(coordinator, entry, description, locale)
for locale in GEOGRAPHY_SENSOR_LOCALES for locale in GEOGRAPHY_SENSOR_LOCALES
for description in GEOGRAPHY_SENSOR_DESCRIPTIONS for description in GEOGRAPHY_SENSOR_DESCRIPTIONS
] ]
else: else:
sensors = [ sensors = [
AirVisualNodeProSensor(coordinator, description) AirVisualNodeProSensor(coordinator, entry, description)
for description in NODE_PRO_SENSOR_DESCRIPTIONS for description in NODE_PRO_SENSOR_DESCRIPTIONS
] ]
@ -221,23 +219,22 @@ class AirVisualGeographySensor(AirVisualEntity, SensorEntity):
def __init__( def __init__(
self, self,
coordinator: DataUpdateCoordinator, coordinator: DataUpdateCoordinator,
config_entry: ConfigEntry, entry: ConfigEntry,
description: SensorEntityDescription, description: SensorEntityDescription,
locale: str, locale: str,
) -> None: ) -> None:
"""Initialize.""" """Initialize."""
super().__init__(coordinator, description) super().__init__(coordinator, entry, description)
self._attr_extra_state_attributes.update( self._attr_extra_state_attributes.update(
{ {
ATTR_CITY: config_entry.data.get(CONF_CITY), ATTR_CITY: entry.data.get(CONF_CITY),
ATTR_STATE: config_entry.data.get(CONF_STATE), ATTR_STATE: entry.data.get(CONF_STATE),
ATTR_COUNTRY: config_entry.data.get(CONF_COUNTRY), ATTR_COUNTRY: entry.data.get(CONF_COUNTRY),
} }
) )
self._attr_name = f"{GEOGRAPHY_SENSOR_LOCALES[locale]} {description.name}" self._attr_name = f"{GEOGRAPHY_SENSOR_LOCALES[locale]} {description.name}"
self._attr_unique_id = f"{config_entry.unique_id}_{locale}_{description.key}" self._attr_unique_id = f"{entry.unique_id}_{locale}_{description.key}"
self._config_entry = config_entry
self._locale = locale self._locale = locale
@property @property
@ -279,16 +276,16 @@ class AirVisualGeographySensor(AirVisualEntity, SensorEntity):
# #
# We use any coordinates in the config entry and, in the case of a geography by # 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: # name, we fall back to the latitude longitude provided in the coordinator data:
latitude = self._config_entry.data.get( latitude = self._entry.data.get(
CONF_LATITUDE, CONF_LATITUDE,
self.coordinator.data["location"]["coordinates"][1], self.coordinator.data["location"]["coordinates"][1],
) )
longitude = self._config_entry.data.get( longitude = self._entry.data.get(
CONF_LONGITUDE, CONF_LONGITUDE,
self.coordinator.data["location"]["coordinates"][0], self.coordinator.data["location"]["coordinates"][0],
) )
if self._config_entry.options[CONF_SHOW_ON_MAP]: if self._entry.options[CONF_SHOW_ON_MAP]:
self._attr_extra_state_attributes[ATTR_LATITUDE] = latitude self._attr_extra_state_attributes[ATTR_LATITUDE] = latitude
self._attr_extra_state_attributes[ATTR_LONGITUDE] = longitude self._attr_extra_state_attributes[ATTR_LONGITUDE] = longitude
self._attr_extra_state_attributes.pop("lati", None) self._attr_extra_state_attributes.pop("lati", None)
@ -304,10 +301,13 @@ class AirVisualNodeProSensor(AirVisualEntity, SensorEntity):
"""Define an AirVisual sensor related to a Node/Pro unit.""" """Define an AirVisual sensor related to a Node/Pro unit."""
def __init__( def __init__(
self, coordinator: DataUpdateCoordinator, description: SensorEntityDescription self,
coordinator: DataUpdateCoordinator,
entry: ConfigEntry,
description: SensorEntityDescription,
) -> None: ) -> None:
"""Initialize.""" """Initialize."""
super().__init__(coordinator, description) super().__init__(coordinator, entry, description)
self._attr_name = ( self._attr_name = (
f"{coordinator.data['settings']['node_name']} Node/Pro: {description.name}" f"{coordinator.data['settings']['node_name']} Node/Pro: {description.name}"

View File

@ -177,10 +177,8 @@ async def test_migration(hass):
], ],
} }
config_entry = MockConfigEntry( entry = MockConfigEntry(domain=DOMAIN, version=1, unique_id="abcde12345", data=conf)
domain=DOMAIN, version=1, unique_id="abcde12345", data=conf entry.add_to_hass(hass)
)
config_entry.add_to_hass(hass)
assert len(hass.config_entries.async_entries(DOMAIN)) == 1 assert len(hass.config_entries.async_entries(DOMAIN)) == 1
@ -222,19 +220,19 @@ async def test_options_flow(hass):
CONF_LONGITUDE: -0.3817765, CONF_LONGITUDE: -0.3817765,
} }
config_entry = MockConfigEntry( entry = MockConfigEntry(
domain=DOMAIN, domain=DOMAIN,
unique_id="51.528308, -0.3817765", unique_id="51.528308, -0.3817765",
data=geography_conf, data=geography_conf,
options={CONF_SHOW_ON_MAP: True}, options={CONF_SHOW_ON_MAP: True},
) )
config_entry.add_to_hass(hass) entry.add_to_hass(hass)
with patch( with patch(
"homeassistant.components.airvisual.async_setup_entry", return_value=True "homeassistant.components.airvisual.async_setup_entry", return_value=True
): ):
await hass.config_entries.async_setup(config_entry.entry_id) await hass.config_entries.async_setup(entry.entry_id)
result = await hass.config_entries.options.async_init(config_entry.entry_id) result = await hass.config_entries.options.async_init(entry.entry_id)
assert result["type"] == data_entry_flow.RESULT_TYPE_FORM assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
assert result["step_id"] == "init" assert result["step_id"] == "init"
@ -244,7 +242,7 @@ async def test_options_flow(hass):
) )
assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
assert config_entry.options == {CONF_SHOW_ON_MAP: False} assert entry.options == {CONF_SHOW_ON_MAP: False}
async def test_step_geography_by_coords(hass): async def test_step_geography_by_coords(hass):