Add restore state to Geofency (#24268)

* Add restore state to Geofency

* Lint

homeassistant/components/geofency/device_tracker.py

@@ -1,12 +1,18 @@
 """Support for the Geofency device tracker platform."""
 import logging
 
+from homeassistant.const import (
+    ATTR_LATITUDE,
+    ATTR_LONGITUDE,
+)
 from homeassistant.core import callback
 from homeassistant.components.device_tracker import SOURCE_TYPE_GPS
 from homeassistant.components.device_tracker.config_entry import (
     DeviceTrackerEntity
 )
 from homeassistant.helpers.dispatcher import async_dispatcher_connect
+from homeassistant.helpers.restore_state import RestoreEntity
+from homeassistant.helpers import device_registry
 
 from . import DOMAIN as GF_DOMAIN, TRACKER_UPDATE
 
@@ -30,15 +36,28 @@ async def async_setup_entry(hass, config_entry, async_add_entities):
     hass.data[GF_DOMAIN]['unsub_device_tracker'][config_entry.entry_id] = \
         async_dispatcher_connect(hass, TRACKER_UPDATE, _receive_data)
 
+    # Restore previously loaded devices
+    dev_reg = await device_registry.async_get_registry(hass)
+    dev_ids = {
+        identifier[1]
+        for device in dev_reg.devices.values()
+        for identifier in device.identifiers
+        if identifier[0] == GF_DOMAIN
+    }
+
+    if dev_ids:
+        hass.data[GF_DOMAIN]['devices'].update(dev_ids)
+        async_add_entities(GeofencyEntity(dev_id) for dev_id in dev_ids)
+
     return True
 
 
-class GeofencyEntity(DeviceTrackerEntity):
+class GeofencyEntity(DeviceTrackerEntity, RestoreEntity):
     """Represent a tracked device."""
 
-    def __init__(self, device, gps, location_name, attributes):
+    def __init__(self, device, gps=None, location_name=None, attributes=None):
         """Set up Geofency entity."""
-        self._attributes = attributes
+        self._attributes = attributes or {}
         self._name = device
         self._location_name = location_name
         self._gps = gps
@@ -95,12 +114,27 @@ class GeofencyEntity(DeviceTrackerEntity):
 
     async def async_added_to_hass(self):
         """Register state update callback."""
+        await super().async_added_to_hass()
         self._unsub_dispatcher = async_dispatcher_connect(
             self.hass, TRACKER_UPDATE, self._async_receive_data)
 
+        if self._attributes:
+            return
+
+        state = await self.async_get_last_state()
+
+        if state is None:
+            self._gps = (None, None)
+            return
+
+        attr = state.attributes
+        self._gps = (attr.get(ATTR_LATITUDE), attr.get(ATTR_LONGITUDE))
+
     async def async_will_remove_from_hass(self):
         """Clean up after entity before removal."""
+        await super().async_will_remove_from_hass()
         self._unsub_dispatcher()
+        self.hass.data[GF_DOMAIN]['devices'].remove(self._unique_id)
 
     @callback
     def _async_receive_data(self, device, gps, location_name, attributes):

tests/components/geofency/test_init.py

@@ -5,13 +5,12 @@ from unittest.mock import patch, Mock
 import pytest
 
 from homeassistant import data_entry_flow
-from homeassistant.components import zone, geofency
+from homeassistant.components import zone
 from homeassistant.components.geofency import (
-    CONF_MOBILE_BEACONS, DOMAIN, TRACKER_UPDATE)
+    CONF_MOBILE_BEACONS, DOMAIN)
 from homeassistant.const import (
     HTTP_OK, HTTP_UNPROCESSABLE_ENTITY, STATE_HOME,
     STATE_NOT_HOME)
-from homeassistant.helpers.dispatcher import DATA_DISPATCHER
 from homeassistant.setup import async_setup_component
 from homeassistant.util import slugify
 
@@ -291,9 +290,6 @@ async def test_beacon_enter_and_exit_car(hass, geofency_client, webhook_id):
     assert STATE_HOME == state_name
 
 
-@pytest.mark.xfail(
-    reason='The device_tracker component does not support unloading yet.'
-)
 async def test_load_unload_entry(hass, geofency_client, webhook_id):
     """Test that the appropriate dispatch signals are added and removed."""
     url = '/api/webhook/{}'.format(webhook_id)
@@ -303,13 +299,23 @@ async def test_load_unload_entry(hass, geofency_client, webhook_id):
     await hass.async_block_till_done()
     assert req.status == HTTP_OK
     device_name = slugify(GPS_ENTER_HOME['device'])
-    state_name = hass.states.get('{}.{}'.format(
-        'device_tracker', device_name)).state
-    assert STATE_HOME == state_name
-    assert len(hass.data[DATA_DISPATCHER][TRACKER_UPDATE]) == 1
+    state_1 = hass.states.get('{}.{}'.format('device_tracker', device_name))
+    assert STATE_HOME == state_1.state
 
+    assert len(hass.data[DOMAIN]['devices']) == 1
     entry = hass.config_entries.async_entries(DOMAIN)[0]
 
-    assert await geofency.async_unload_entry(hass, entry)
+    assert await hass.config_entries.async_unload(entry.entry_id)
     await hass.async_block_till_done()
-    assert not hass.data[DATA_DISPATCHER][TRACKER_UPDATE]
+    assert len(hass.data[DOMAIN]['devices']) == 0
+
+    assert await hass.config_entries.async_setup(entry.entry_id)
+    await hass.async_block_till_done()
+
+    state_2 = hass.states.get('{}.{}'.format('device_tracker', device_name))
+    assert state_2 is not None
+    assert state_1 is not state_2
+
+    assert STATE_HOME == state_2.state
+    assert state_2.attributes['latitude'] == HOME_LATITUDE
+    assert state_2.attributes['longitude'] == HOME_LONGITUDE