GeoNet NZ Quakes code improvements (home-assistant#32338)

* code quality improvements

* code quality improvements and fixed tests

* explicitly set unique ids

* improve unique id creation

* remove entities from entity registry

* added test for removing entities from entity registry

* revert entity registry handling from sensor and test code

* check for entity registry removal in geolocation test case

* make import absolute; isort

* change quality scale

homeassistant/components/geonetnz_quakes/__init__.py

@@ -12,7 +12,6 @@
     CONF_LONGITUDE,
     CONF_RADIUS,
     CONF_SCAN_INTERVAL,
-    CONF_UNIT_SYSTEM,
     CONF_UNIT_SYSTEM_IMPERIAL,
     LENGTH_MILES,
 )
@@ -22,7 +21,6 @@
 from homeassistant.helpers.event import async_track_time_interval
 from homeassistant.util.unit_system import METRIC_SYSTEM
 
-from .config_flow import configured_instances
 from .const import (
     CONF_MINIMUM_MAGNITUDE,
     CONF_MMI,
@@ -42,8 +40,8 @@
     {
         DOMAIN: vol.Schema(
             {
-                vol.Optional(CONF_LATITUDE): cv.latitude,
-                vol.Optional(CONF_LONGITUDE): cv.longitude,
+                vol.Inclusive(CONF_LATITUDE, "coordinates"): cv.latitude,
+                vol.Inclusive(CONF_LONGITUDE, "coordinates"): cv.longitude,
                 vol.Optional(CONF_MMI, default=DEFAULT_MMI): vol.All(
                     vol.Coerce(int), vol.Range(min=-1, max=8)
                 ),
@@ -67,16 +65,11 @@ async def async_setup(hass, config):
         return True
 
     conf = config[DOMAIN]
-
     latitude = conf.get(CONF_LATITUDE, hass.config.latitude)
     longitude = conf.get(CONF_LONGITUDE, hass.config.longitude)
     mmi = conf[CONF_MMI]
     scan_interval = conf[CONF_SCAN_INTERVAL]
 
-    identifier = f"{latitude}, {longitude}"
-    if identifier in configured_instances(hass):
-        return True
-
     hass.async_create_task(
         hass.config_entries.flow.async_init(
             DOMAIN,
@@ -97,18 +90,15 @@ async def async_setup(hass, config):
 
 async def async_setup_entry(hass, config_entry):
     """Set up the GeoNet NZ Quakes component as config entry."""
-    if DOMAIN not in hass.data:
-        hass.data[DOMAIN] = {}
-    if FEED not in hass.data[DOMAIN]:
-        hass.data[DOMAIN][FEED] = {}
+    hass.data.setdefault(DOMAIN, {})
+    feeds = hass.data[DOMAIN].setdefault(FEED, {})
 
     radius = config_entry.data[CONF_RADIUS]
-    unit_system = config_entry.data[CONF_UNIT_SYSTEM]
-    if unit_system == CONF_UNIT_SYSTEM_IMPERIAL:
+    if hass.config.units.name == CONF_UNIT_SYSTEM_IMPERIAL:
         radius = METRIC_SYSTEM.length(radius, LENGTH_MILES)
     # Create feed entity manager for all platforms.
-    manager = GeonetnzQuakesFeedEntityManager(hass, config_entry, radius, unit_system)
-    hass.data[DOMAIN][FEED][config_entry.entry_id] = manager
+    manager = GeonetnzQuakesFeedEntityManager(hass, config_entry, radius)
+    feeds[config_entry.entry_id] = manager
     _LOGGER.debug("Feed entity manager added for %s", config_entry.entry_id)
     await manager.async_init()
     return True
@@ -130,7 +120,7 @@ async def async_unload_entry(hass, config_entry):
 class GeonetnzQuakesFeedEntityManager:
     """Feed Entity Manager for GeoNet NZ Quakes feed."""
 
-    def __init__(self, hass, config_entry, radius_in_km, unit_system):
+    def __init__(self, hass, config_entry, radius_in_km):
         """Initialize the Feed Entity Manager."""
         self._hass = hass
         self._config_entry = config_entry
@@ -153,7 +143,6 @@ def __init__(self, hass, config_entry, radius_in_km, unit_system):
         )
         self._config_entry_id = config_entry.entry_id
         self._scan_interval = timedelta(seconds=config_entry.data[CONF_SCAN_INTERVAL])
-        self._unit_system = unit_system
         self._track_time_remove_callback = None
         self._status_info = None
         self.listeners = []
@@ -212,8 +201,8 @@ async def _generate_entity(self, external_id):
             self._hass,
             self.async_event_new_entity(),
             self,
+            self._config_entry.unique_id,
             external_id,
-            self._unit_system,
         )
 
     async def _update_entity(self, external_id):

homeassistant/components/geonetnz_quakes/config_flow.py

@@ -9,14 +9,10 @@
     CONF_LONGITUDE,
     CONF_RADIUS,
     CONF_SCAN_INTERVAL,
-    CONF_UNIT_SYSTEM,
-    CONF_UNIT_SYSTEM_IMPERIAL,
-    CONF_UNIT_SYSTEM_METRIC,
 )
-from homeassistant.core import callback
 from homeassistant.helpers import config_validation as cv
 
-from .const import (
+from .const import (  # pylint: disable=unused-import
     CONF_MINIMUM_MAGNITUDE,
     CONF_MMI,
     DEFAULT_MINIMUM_MAGNITUDE,
@@ -26,37 +22,27 @@
     DOMAIN,
 )
 
-_LOGGER = logging.getLogger(__name__)
-
+DATA_SCHEMA = vol.Schema(
+    {
+        vol.Optional(CONF_MMI, default=DEFAULT_MMI): vol.All(
+            vol.Coerce(int), vol.Range(min=-1, max=8)
+        ),
+        vol.Optional(CONF_RADIUS, default=DEFAULT_RADIUS): cv.positive_int,
+    }
+)
 
-@callback
-def configured_instances(hass):
-    """Return a set of configured GeoNet NZ Quakes instances."""
-    return set(
-        f"{entry.data[CONF_LATITUDE]}, {entry.data[CONF_LONGITUDE]}"
-        for entry in hass.config_entries.async_entries(DOMAIN)
-    )
+_LOGGER = logging.getLogger(__name__)
 
 
-@config_entries.HANDLERS.register(DOMAIN)
-class GeonetnzQuakesFlowHandler(config_entries.ConfigFlow):
+class GeonetnzQuakesFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
     """Handle a GeoNet NZ Quakes config flow."""
 
     CONNECTION_CLASS = config_entries.CONN_CLASS_CLOUD_POLL
 
     async def _show_form(self, errors=None):
         """Show the form to the user."""
-        data_schema = vol.Schema(
-            {
-                vol.Optional(CONF_MMI, default=DEFAULT_MMI): vol.All(
-                    vol.Coerce(int), vol.Range(min=-1, max=8)
-                ),
-                vol.Optional(CONF_RADIUS, default=DEFAULT_RADIUS): cv.positive_int,
-            }
-        )
-
         return self.async_show_form(
-            step_id="user", data_schema=data_schema, errors=errors or {}
+            step_id="user", data_schema=DATA_SCHEMA, errors=errors or {}
         )
 
     async def async_step_import(self, import_config):
@@ -75,13 +61,9 @@ async def async_step_user(self, user_input=None):
         user_input[CONF_LONGITUDE] = longitude
 
         identifier = f"{user_input[CONF_LATITUDE]}, {user_input[CONF_LONGITUDE]}"
-        if identifier in configured_instances(self.hass):
-            return await self._show_form({"base": "identifier_exists"})
 
-        if self.hass.config.units.name == CONF_UNIT_SYSTEM_IMPERIAL:
-            user_input[CONF_UNIT_SYSTEM] = CONF_UNIT_SYSTEM_IMPERIAL
-        else:
-            user_input[CONF_UNIT_SYSTEM] = CONF_UNIT_SYSTEM_METRIC
+        await self.async_set_unique_id(identifier)
+        self._abort_if_unique_id_configured()
 
         scan_interval = user_input.get(CONF_SCAN_INTERVAL, DEFAULT_SCAN_INTERVAL)
         user_input[CONF_SCAN_INTERVAL] = scan_interval.seconds

homeassistant/components/geonetnz_quakes/geo_location.py

@@ -12,6 +12,7 @@
 )
 from homeassistant.core import callback
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.entity_registry import async_get_registry
 from homeassistant.util.unit_system import IMPERIAL_SYSTEM
 
 from .const import DOMAIN, FEED
@@ -26,6 +27,9 @@
 ATTR_PUBLICATION_DATE = "publication_date"
 ATTR_QUALITY = "quality"
 
+# An update of this entity is not making a web request, but uses internal data only.
+PARALLEL_UPDATES = 0
+
 SOURCE = "geonetnz_quakes"
 
 
@@ -34,9 +38,9 @@ async def async_setup_entry(hass, entry, async_add_entities):
     manager = hass.data[DOMAIN][FEED][entry.entry_id]
 
     @callback
-    def async_add_geolocation(feed_manager, external_id, unit_system):
+    def async_add_geolocation(feed_manager, integration_id, external_id):
         """Add gelocation entity from feed."""
-        new_entity = GeonetnzQuakesEvent(feed_manager, external_id, unit_system)
+        new_entity = GeonetnzQuakesEvent(feed_manager, integration_id, external_id)
         _LOGGER.debug("Adding geolocation %s", new_entity)
         async_add_entities([new_entity], True)
 
@@ -45,18 +49,20 @@ def async_add_geolocation(feed_manager, external_id, unit_system):
             hass, manager.async_event_new_entity(), async_add_geolocation
         )
     )
+    # Do not wait for update here so that the setup can be completed and because an
+    # update will fetch data from the feed via HTTP and then process that data.
     hass.async_create_task(manager.async_update())
     _LOGGER.debug("Geolocation setup done")
 
 
 class GeonetnzQuakesEvent(GeolocationEvent):
     """This represents an external event with GeoNet NZ Quakes feed data."""
 
-    def __init__(self, feed_manager, external_id, unit_system):
+    def __init__(self, feed_manager, integration_id, external_id):
         """Initialize entity with data from feed entry."""
         self._feed_manager = feed_manager
+        self._integration_id = integration_id
         self._external_id = external_id
-        self._unit_system = unit_system
         self._title = None
         self._distance = None
         self._latitude = None
@@ -88,6 +94,9 @@ async def async_will_remove_from_hass(self) -> None:
         """Call when entity will be removed from hass."""
         self._remove_signal_delete()
         self._remove_signal_update()
+        # Remove from entity registry.
+        entity_registry = await async_get_registry(self.hass)
+        entity_registry.async_remove(self.entity_id)
 
     @callback
     def _delete_callback(self):
@@ -115,7 +124,7 @@ def _update_from_feed(self, feed_entry):
         """Update the internal state from the provided feed entry."""
         self._title = feed_entry.title
         # Convert distance if not metric system.
-        if self._unit_system == CONF_UNIT_SYSTEM_IMPERIAL:
+        if self.hass.config.units.name == CONF_UNIT_SYSTEM_IMPERIAL:
             self._distance = IMPERIAL_SYSTEM.length(
                 feed_entry.distance_to_home, LENGTH_KILOMETERS
             )
@@ -131,6 +140,11 @@ def _update_from_feed(self, feed_entry):
         self._quality = feed_entry.quality
         self._time = feed_entry.time
 
+    @property
+    def unique_id(self) -> Optional[str]:
+        """Return a unique ID containing latitude/longitude and external id."""
+        return f"{self._integration_id}_{self._external_id}"
+
     @property
     def icon(self):
         """Return the icon to use in the frontend, if any."""
@@ -164,7 +178,7 @@ def longitude(self) -> Optional[float]:
     @property
     def unit_of_measurement(self):
         """Return the unit of measurement."""
-        if self._unit_system == CONF_UNIT_SYSTEM_IMPERIAL:
+        if self.hass.config.units.name == CONF_UNIT_SYSTEM_IMPERIAL:
             return LENGTH_MILES
         return LENGTH_KILOMETERS
 

homeassistant/components/geonetnz_quakes/manifest.json

@@ -5,5 +5,6 @@
   "documentation": "https://www.home-assistant.io/integrations/geonetnz_quakes",
   "requirements": ["aio_geojson_geonetnz_quakes==0.12"],
   "dependencies": [],
-  "codeowners": ["@exxamalte"]
-}
+  "codeowners": ["@exxamalte"],
+  "quality_scale": "platinum"
+}

homeassistant/components/geonetnz_quakes/sensor.py

@@ -22,21 +22,25 @@
 DEFAULT_ICON = "mdi:pulse"
 DEFAULT_UNIT_OF_MEASUREMENT = "quakes"
 
+# An update of this entity is not making a web request, but uses internal data only.
+PARALLEL_UPDATES = 0
+
 
 async def async_setup_entry(hass, entry, async_add_entities):
     """Set up the GeoNet NZ Quakes Feed platform."""
     manager = hass.data[DOMAIN][FEED][entry.entry_id]
-    sensor = GeonetnzQuakesSensor(entry.entry_id, entry.title, manager)
+    sensor = GeonetnzQuakesSensor(entry.entry_id, entry.unique_id, entry.title, manager)
     async_add_entities([sensor])
     _LOGGER.debug("Sensor setup done")
 
 
 class GeonetnzQuakesSensor(Entity):
     """This is a status sensor for the GeoNet NZ Quakes integration."""
 
-    def __init__(self, config_entry_id, config_title, manager):
+    def __init__(self, config_entry_id, config_unique_id, config_title, manager):
         """Initialize entity."""
         self._config_entry_id = config_entry_id
+        self._config_unique_id = config_unique_id
         self._config_title = config_title
         self._manager = manager
         self._status = None
@@ -90,11 +94,10 @@ def _update_from_status_info(self, status_info):
         self._last_update = (
             dt.as_utc(status_info.last_update) if status_info.last_update else None
         )
-        self._last_update_successful = (
-            dt.as_utc(status_info.last_update_successful)
-            if status_info.last_update_successful
-            else None
-        )
+        if status_info.last_update_successful:
+            self._last_update_successful = dt.as_utc(status_info.last_update_successful)
+        else:
+            self._last_update_successful = None
         self._last_timestamp = status_info.last_timestamp
         self._total = status_info.total
         self._created = status_info.created
@@ -106,6 +109,11 @@ def state(self):
         """Return the state of the sensor."""
         return self._total
 
+    @property
+    def unique_id(self) -> str:
+        """Return a unique ID containing latitude/longitude."""
+        return self._config_unique_id
+
     @property
     def name(self) -> Optional[str]:
         """Return the name of the entity."""

homeassistant/components/geonetnz_quakes/strings.json

@@ -10,8 +10,8 @@
         }
       }
     },
-    "error": {
-      "identifier_exists": "Location already registered"
+    "abort": {
+      "already_configured": "Location is already configured."
     }
   }
 }

tests/components/geonetnz_quakes/conftest.py

@@ -0,0 +1,36 @@
+"""Configuration for GeoNet NZ Quakes tests."""
+import pytest
+
+from homeassistant.components.geonetnz_quakes import (
+    CONF_MINIMUM_MAGNITUDE,
+    CONF_MMI,
+    DOMAIN,
+)
+from homeassistant.const import (
+    CONF_LATITUDE,
+    CONF_LONGITUDE,
+    CONF_RADIUS,
+    CONF_SCAN_INTERVAL,
+    CONF_UNIT_SYSTEM,
+)
+
+from tests.common import MockConfigEntry
+
+
+@pytest.fixture
+def config_entry():
+    """Create a mock GeoNet NZ Quakes config entry."""
+    return MockConfigEntry(
+        domain=DOMAIN,
+        data={
+            CONF_LATITUDE: -41.2,
+            CONF_LONGITUDE: 174.7,
+            CONF_RADIUS: 25,
+            CONF_UNIT_SYSTEM: "metric",
+            CONF_SCAN_INTERVAL: 300.0,
+            CONF_MMI: 4,
+            CONF_MINIMUM_MAGNITUDE: 0.0,
+        },
+        title="-41.2, 174.7",
+        unique_id="-41.2, 174.7",
+    )