Use runtime_data and HassKey in geofency (#144886)

homeassistant/components/geofency/__init__.py

@@ -20,9 +20,12 @@ from homeassistant.helpers import config_entry_flow, config_validation as cv
 from homeassistant.helpers.dispatcher import async_dispatcher_send
 from homeassistant.helpers.typing import ConfigType
 from homeassistant.util import slugify
+from homeassistant.util.hass_dict import HassKey
 
 from .const import DOMAIN
 
+type GeofencyConfigEntry = ConfigEntry[set[str]]
+
 PLATFORMS = [Platform.DEVICE_TRACKER]
 
 CONF_MOBILE_BEACONS = "mobile_beacons"
@@ -75,15 +78,13 @@ WEBHOOK_SCHEMA = vol.Schema(
     extra=vol.ALLOW_EXTRA,
 )
 
+_DATA_GEOFENCY: HassKey[list[str]] = HassKey(DOMAIN)
+
 
 async def async_setup(hass: HomeAssistant, hass_config: ConfigType) -> bool:
     """Set up the Geofency component."""
-    config = hass_config.get(DOMAIN, {})
+    mobile_beacons = hass_config.get(DOMAIN, {}).get(CONF_MOBILE_BEACONS, [])
-    mobile_beacons = config.get(CONF_MOBILE_BEACONS, [])
+    hass.data[_DATA_GEOFENCY] = [slugify(beacon) for beacon in mobile_beacons]
-    hass.data[DOMAIN] = {
-        "beacons": [slugify(beacon) for beacon in mobile_beacons],
-        "devices": set(),
-    }
     return True
 
 
@@ -98,7 +99,7 @@ async def handle_webhook(
             text=error.error_message, status=HTTPStatus.UNPROCESSABLE_ENTITY
         )
 
-    if _is_mobile_beacon(data, hass.data[DOMAIN]["beacons"]):
+    if _is_mobile_beacon(data, hass.data[_DATA_GEOFENCY]):
         return _set_location(hass, data, None)
     if data["entry"] == LOCATION_ENTRY:
         location_name = data["name"]
@@ -139,8 +140,9 @@ def _set_location(hass, data, location_name):
     return web.Response(text=f"Setting location for {device}")
 
 
-async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_setup_entry(hass: HomeAssistant, entry: GeofencyConfigEntry) -> bool:
     """Configure based on config entry."""
+    entry.runtime_data = set()
     webhook.async_register(
         hass, DOMAIN, "Geofency", entry.data[CONF_WEBHOOK_ID], handle_webhook
     )
@@ -149,7 +151,7 @@ async def async_setup_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
     return True
 
 
-async def async_unload_entry(hass: HomeAssistant, entry: ConfigEntry) -> bool:
+async def async_unload_entry(hass: HomeAssistant, entry: GeofencyConfigEntry) -> bool:
     """Unload a config entry."""
     webhook.async_unregister(hass, entry.data[CONF_WEBHOOK_ID])
     return await hass.config_entries.async_unload_platforms(entry, PLATFORMS)

homeassistant/components/geofency/device_tracker.py

@@ -1,7 +1,6 @@
 """Support for the Geofency device tracker platform."""
 
 from homeassistant.components.device_tracker import TrackerEntity
-from homeassistant.config_entries import ConfigEntry
 from homeassistant.const import ATTR_LATITUDE, ATTR_LONGITUDE
 from homeassistant.core import HomeAssistant, callback
 from homeassistant.helpers import device_registry as dr
@@ -10,12 +9,13 @@ from homeassistant.helpers.dispatcher import async_dispatcher_connect
 from homeassistant.helpers.entity_platform import AddConfigEntryEntitiesCallback
 from homeassistant.helpers.restore_state import RestoreEntity
 
-from . import DOMAIN, TRACKER_UPDATE
+from . import TRACKER_UPDATE, GeofencyConfigEntry
+from .const import DOMAIN
 
 
 async def async_setup_entry(
     hass: HomeAssistant,
-    config_entry: ConfigEntry,
+    config_entry: GeofencyConfigEntry,
     async_add_entities: AddConfigEntryEntitiesCallback,
 ) -> None:
     """Set up Geofency config entry."""
@@ -23,12 +23,14 @@ async def async_setup_entry(
     @callback
     def _receive_data(device, gps, location_name, attributes):
         """Fire HA event to set location."""
-        if device in hass.data[DOMAIN]["devices"]:
+        if device in config_entry.runtime_data:
             return
 
-        hass.data[DOMAIN]["devices"].add(device)
+        config_entry.runtime_data.add(device)
 
-        async_add_entities([GeofencyEntity(device, gps, location_name, attributes)])
+        async_add_entities(
+            [GeofencyEntity(config_entry, device, gps, location_name, attributes)]
+        )
 
     config_entry.async_on_unload(
         async_dispatcher_connect(hass, TRACKER_UPDATE, _receive_data)
@@ -45,8 +47,8 @@ async def async_setup_entry(
     }
 
     if dev_ids:
-        hass.data[DOMAIN]["devices"].update(dev_ids)
+        config_entry.runtime_data.update(dev_ids)
-        async_add_entities(GeofencyEntity(dev_id) for dev_id in dev_ids)
+        async_add_entities(GeofencyEntity(config_entry, dev_id) for dev_id in dev_ids)
 
 
 class GeofencyEntity(TrackerEntity, RestoreEntity):
@@ -55,8 +57,9 @@ class GeofencyEntity(TrackerEntity, RestoreEntity):
     _attr_has_entity_name = True
     _attr_name = None
 
-    def __init__(self, device, gps=None, location_name=None, attributes=None):
+    def __init__(self, entry, device, gps=None, location_name=None, attributes=None):
         """Set up Geofency entity."""
+        self._entry = entry
         self._attr_extra_state_attributes = attributes or {}
         self._name = device
         self._attr_location_name = location_name
@@ -93,7 +96,7 @@ class GeofencyEntity(TrackerEntity, RestoreEntity):
         """Clean up after entity before removal."""
         await super().async_will_remove_from_hass()
         self._unsub_dispatcher()
-        self.hass.data[DOMAIN]["devices"].remove(self.unique_id)
+        self._entry.runtime_data.remove(self.unique_id)
 
     @callback
     def _async_receive_data(self, device, gps, location_name, attributes):

tests/components/geofency/test_init.py

@@ -318,12 +318,11 @@ async def test_load_unload_entry(
     state_1 = hass.states.get(f"device_tracker.{device_name}")
     assert state_1.state == STATE_HOME
 
-    assert len(hass.data[DOMAIN]["devices"]) == 1
     entry = hass.config_entries.async_entries(DOMAIN)[0]
+    assert len(entry.runtime_data) == 1
 
     assert await hass.config_entries.async_unload(entry.entry_id)
     await hass.async_block_till_done()
-    assert len(hass.data[DOMAIN]["devices"]) == 0
 
     assert await hass.config_entries.async_setup(entry.entry_id)
     await hass.async_block_till_done()