GeoNet NZ Quakes code improvements (#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 @@ from homeassistant.const import (
     CONF_LONGITUDE,
     CONF_RADIUS,
     CONF_SCAN_INTERVAL,
-    CONF_UNIT_SYSTEM,
     CONF_UNIT_SYSTEM_IMPERIAL,
     LENGTH_MILES,
 )
@@ -22,7 +21,6 @@ from homeassistant.helpers.dispatcher import async_dispatcher_send
 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 @@ CONFIG_SCHEMA = vol.Schema(
     {
         DOMAIN: vol.Schema(
             {
-                vol.Optional(CONF_LATITUDE): cv.latitude,
+                vol.Inclusive(CONF_LATITUDE, "coordinates"): cv.latitude,
-                vol.Optional(CONF_LONGITUDE): cv.longitude,
+                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.setdefault(DOMAIN, {})
-        hass.data[DOMAIN] = {}
+    feeds = hass.data[DOMAIN].setdefault(FEED, {})
-    if FEED not in hass.data[DOMAIN]:
-        hass.data[DOMAIN][FEED] = {}
 
     radius = config_entry.data[CONF_RADIUS]
-    unit_system = config_entry.data[CONF_UNIT_SYSTEM]
+    if hass.config.units.name == CONF_UNIT_SYSTEM_IMPERIAL:
-    if unit_system == 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)
+    manager = GeonetnzQuakesFeedEntityManager(hass, config_entry, radius)
-    hass.data[DOMAIN][FEED][config_entry.entry_id] = manager
+    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 @@ class GeonetnzQuakesFeedEntityManager:
         )
         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 @@ class GeonetnzQuakesFeedEntityManager:
             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 @@ from homeassistant.const import (
     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 @@ from .const import (
     DOMAIN,
 )
 
+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,
+    }
+)
+
 _LOGGER = logging.getLogger(__name__)
 
 
-@callback
+class GeonetnzQuakesFlowHandler(config_entries.ConfigFlow, domain=DOMAIN):
-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)
-    )
-
-
-@config_entries.HANDLERS.register(DOMAIN)
-class GeonetnzQuakesFlowHandler(config_entries.ConfigFlow):
     """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 @@ class GeonetnzQuakesFlowHandler(config_entries.ConfigFlow):
         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:
+        await self.async_set_unique_id(identifier)
-            user_input[CONF_UNIT_SYSTEM] = CONF_UNIT_SYSTEM_IMPERIAL
+        self._abort_if_unique_id_configured()
-        else:
-            user_input[CONF_UNIT_SYSTEM] = CONF_UNIT_SYSTEM_METRIC
 
         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.const import (
 )
 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_MMI = "mmi"
 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,6 +49,8 @@ async def async_setup_entry(hass, entry, async_add_entities):
             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")
 
@@ -52,11 +58,11 @@ async def async_setup_entry(hass, entry, async_add_entities):
 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 @@ class GeonetnzQuakesEvent(GeolocationEvent):
         """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 @@ class GeonetnzQuakesEvent(GeolocationEvent):
         """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 @@ class GeonetnzQuakesEvent(GeolocationEvent):
         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 @@ class GeonetnzQuakesEvent(GeolocationEvent):
     @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,11 +22,14 @@ ATTR_REMOVED = "removed"
 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")
 
@@ -34,9 +37,10 @@ async def async_setup_entry(hass, entry, async_add_entities):
 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 @@ class GeonetnzQuakesSensor(Entity):
         self._last_update = (
             dt.as_utc(status_info.last_update) if status_info.last_update else None
         )
-        self._last_update_successful = (
+        if status_info.last_update_successful:
-            dt.as_utc(status_info.last_update_successful)
+            self._last_update_successful = dt.as_utc(status_info.last_update_successful)
-            if status_info.last_update_successful
+        else:
-            else None
+            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 @@ class GeonetnzQuakesSensor(Entity):
         """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": {
+    "abort": {
-      "identifier_exists": "Location already registered"
+      "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",
+    )

tests/components/geonetnz_quakes/test_config_flow.py

@@ -1,18 +1,11 @@
 """Define tests for the GeoNet NZ Quakes config flow."""
 from datetime import timedelta
 
-from asynctest import CoroutineMock, patch
-import pytest
-
 from homeassistant import data_entry_flow
 from homeassistant.components.geonetnz_quakes import (
     CONF_MINIMUM_MAGNITUDE,
     CONF_MMI,
     DOMAIN,
-    FEED,
-    async_setup_entry,
-    async_unload_entry,
-    config_flow,
 )
 from homeassistant.const import (
     CONF_LATITUDE,
@@ -22,46 +15,24 @@ from homeassistant.const import (
     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",
-    )
-
 
 async def test_duplicate_error(hass, config_entry):
     """Test that errors are shown when duplicates are added."""
     conf = {CONF_LATITUDE: -41.2, CONF_LONGITUDE: 174.7, CONF_RADIUS: 25}
-
     config_entry.add_to_hass(hass)
-    flow = config_flow.GeonetnzQuakesFlowHandler()
-    flow.hass = hass
 
-    result = await flow.async_step_user(user_input=conf)
+    result = await hass.config_entries.flow.async_init(
-    assert result["errors"] == {"base": "identifier_exists"}
+        DOMAIN, context={"source": "user"}, data=conf
+    )
+    assert result["type"] == data_entry_flow.RESULT_TYPE_ABORT
+    assert result["reason"] == "already_configured"
 
 
 async def test_show_form(hass):
     """Test that the form is served with no input."""
-    flow = config_flow.GeonetnzQuakesFlowHandler()
+    result = await hass.config_entries.flow.async_init(
-    flow.hass = hass
+        DOMAIN, context={"source": "user"}
-
+    )
-    result = await flow.async_step_user(user_input=None)
-
     assert result["type"] == data_entry_flow.RESULT_TYPE_FORM
     assert result["step_id"] == "user"
 
@@ -78,10 +49,9 @@ async def test_step_import(hass):
         CONF_MINIMUM_MAGNITUDE: 2.5,
     }
 
-    flow = config_flow.GeonetnzQuakesFlowHandler()
+    result = await hass.config_entries.flow.async_init(
-    flow.hass = hass
+        DOMAIN, context={"source": "import"}, data=conf
-
+    )
-    result = await flow.async_step_import(import_config=conf)
     assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
     assert result["title"] == "-41.2, 174.7"
     assert result["data"] == {
@@ -101,10 +71,9 @@ async def test_step_user(hass):
     hass.config.longitude = 174.7
     conf = {CONF_RADIUS: 25, CONF_MMI: 4}
 
-    flow = config_flow.GeonetnzQuakesFlowHandler()
+    result = await hass.config_entries.flow.async_init(
-    flow.hass = hass
+        DOMAIN, context={"source": "user"}, data=conf
-
+    )
-    result = await flow.async_step_user(user_input=conf)
     assert result["type"] == data_entry_flow.RESULT_TYPE_CREATE_ENTRY
     assert result["title"] == "-41.2, 174.7"
     assert result["data"] == {
@@ -112,25 +81,6 @@ async def test_step_user(hass):
         CONF_LONGITUDE: 174.7,
         CONF_RADIUS: 25,
         CONF_MMI: 4,
-        CONF_UNIT_SYSTEM: "metric",
         CONF_SCAN_INTERVAL: 300.0,
         CONF_MINIMUM_MAGNITUDE: 0.0,
     }
-
-
-async def test_component_unload_config_entry(hass, config_entry):
-    """Test that loading and unloading of a config entry works."""
-    config_entry.add_to_hass(hass)
-    with patch(
-        "aio_geojson_geonetnz_quakes.GeonetnzQuakesFeedManager.update",
-        new_callable=CoroutineMock,
-    ) as mock_feed_manager_update:
-        # Load config entry.
-        assert await async_setup_entry(hass, config_entry)
-        await hass.async_block_till_done()
-        assert mock_feed_manager_update.call_count == 1
-        assert hass.data[DOMAIN][FEED][config_entry.entry_id] is not None
-        # Unload config entry.
-        assert await async_unload_entry(hass, config_entry)
-        await hass.async_block_till_done()
-        assert hass.data[DOMAIN][FEED].get(config_entry.entry_id) is None

tests/components/geonetnz_quakes/test_geo_location.py

@@ -1,11 +1,11 @@
 """The tests for the GeoNet NZ Quakes Feed integration."""
 import datetime
 
-from asynctest import CoroutineMock, patch
+from asynctest import patch
 
 from homeassistant.components import geonetnz_quakes
 from homeassistant.components.geo_location import ATTR_SOURCE
-from homeassistant.components.geonetnz_quakes import DEFAULT_SCAN_INTERVAL
+from homeassistant.components.geonetnz_quakes import DEFAULT_SCAN_INTERVAL, DOMAIN, FEED
 from homeassistant.components.geonetnz_quakes.geo_location import (
     ATTR_DEPTH,
     ATTR_EXTERNAL_ID,
@@ -25,6 +25,7 @@ from homeassistant.const import (
     CONF_RADIUS,
     EVENT_HOMEASSISTANT_START,
 )
+from homeassistant.helpers.entity_registry import async_get_registry
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 from homeassistant.util.unit_system import IMPERIAL_SYSTEM
@@ -62,7 +63,7 @@ async def test_setup(hass):
     # Patching 'utcnow' to gain more control over the timed update.
     utcnow = dt_util.utcnow()
     with patch("homeassistant.util.dt.utcnow", return_value=utcnow), patch(
-        "aio_geojson_client.feed.GeoJsonFeed.update", new_callable=CoroutineMock
+        "aio_geojson_client.feed.GeoJsonFeed.update"
     ) as mock_feed_update:
         mock_feed_update.return_value = "OK", [mock_entry_1, mock_entry_2, mock_entry_3]
         assert await async_setup_component(hass, geonetnz_quakes.DOMAIN, CONFIG)
@@ -73,6 +74,8 @@ async def test_setup(hass):
         all_states = hass.states.async_all()
         # 3 geolocation and 1 sensor entities
         assert len(all_states) == 4
+        entity_registry = await async_get_registry(hass)
+        assert len(entity_registry.entities) == 4
 
         state = hass.states.get("geo_location.title_1")
         assert state is not None
@@ -151,6 +154,7 @@ async def test_setup(hass):
 
         all_states = hass.states.async_all()
         assert len(all_states) == 1
+        assert len(entity_registry.entities) == 1
 
 
 async def test_setup_imperial(hass):
@@ -162,15 +166,9 @@ async def test_setup_imperial(hass):
     # Patching 'utcnow' to gain more control over the timed update.
     utcnow = dt_util.utcnow()
     with patch("homeassistant.util.dt.utcnow", return_value=utcnow), patch(
-        "aio_geojson_client.feed.GeoJsonFeed.update", new_callable=CoroutineMock
+        "aio_geojson_client.feed.GeoJsonFeed.update"
     ) as mock_feed_update, patch(
-        "aio_geojson_client.feed.GeoJsonFeed.__init__",
+        "aio_geojson_client.feed.GeoJsonFeed.last_timestamp", create=True
-        new_callable=CoroutineMock,
-        create=True,
-    ) as mock_feed_init, patch(
-        "aio_geojson_client.feed.GeoJsonFeed.last_timestamp",
-        new_callable=CoroutineMock,
-        create=True,
     ):
         mock_feed_update.return_value = "OK", [mock_entry_1]
         assert await async_setup_component(hass, geonetnz_quakes.DOMAIN, CONFIG)
@@ -182,7 +180,12 @@ async def test_setup_imperial(hass):
         assert len(all_states) == 2
 
         # Test conversion of 200 miles to kilometers.
-        assert mock_feed_init.call_args[1].get("filter_radius") == 321.8688
+        feeds = hass.data[DOMAIN][FEED]
+        assert feeds is not None
+        assert len(feeds) == 1
+        manager = list(feeds.values())[0]
+        # Ensure that the filter value in km is correctly set.
+        assert manager._feed_manager._feed._filter_radius == 321.8688
 
         state = hass.states.get("geo_location.title_1")
         assert state is not None
@@ -196,4 +199,5 @@ async def test_setup_imperial(hass):
             ATTR_SOURCE: "geonetnz_quakes",
             ATTR_ICON: "mdi:pulse",
         }
+        # 15.5km (as defined in mock entry) has been converted to 9.6mi.
         assert float(state.state) == 9.6

tests/components/geonetnz_quakes/test_init.py

@@ -0,0 +1,21 @@
+"""Define tests for the GeoNet NZ Quakes general setup."""
+from asynctest import patch
+
+from homeassistant.components.geonetnz_quakes import DOMAIN, FEED
+
+
+async def test_component_unload_config_entry(hass, config_entry):
+    """Test that loading and unloading of a config entry works."""
+    config_entry.add_to_hass(hass)
+    with patch(
+        "aio_geojson_geonetnz_quakes.GeonetnzQuakesFeedManager.update"
+    ) as mock_feed_manager_update:
+        # Load config entry.
+        assert await hass.config_entries.async_setup(config_entry.entry_id)
+        await hass.async_block_till_done()
+        assert mock_feed_manager_update.call_count == 1
+        assert hass.data[DOMAIN][FEED][config_entry.entry_id] is not None
+        # Unload config entry.
+        assert await hass.config_entries.async_unload(config_entry.entry_id)
+        await hass.async_block_till_done()
+        assert hass.data[DOMAIN][FEED].get(config_entry.entry_id) is None

tests/components/geonetnz_quakes/test_sensor.py

@@ -1,7 +1,7 @@
 """The tests for the GeoNet NZ Quakes Feed integration."""
 import datetime
 
-from asynctest import CoroutineMock, patch
+from asynctest import patch
 
 from homeassistant.components import geonetnz_quakes
 from homeassistant.components.geonetnz_quakes import DEFAULT_SCAN_INTERVAL
@@ -55,7 +55,7 @@ async def test_setup(hass):
     # Patching 'utcnow' to gain more control over the timed update.
     utcnow = dt_util.utcnow()
     with patch("homeassistant.util.dt.utcnow", return_value=utcnow), patch(
-        "aio_geojson_client.feed.GeoJsonFeed.update", new_callable=CoroutineMock
+        "aio_geojson_client.feed.GeoJsonFeed.update"
     ) as mock_feed_update:
         mock_feed_update.return_value = "OK", [mock_entry_1, mock_entry_2, mock_entry_3]
         assert await async_setup_component(hass, geonetnz_quakes.DOMAIN, CONFIG)