From 5a9970e63ccad5ee84e73593a0edcdc4509833a6 Mon Sep 17 00:00:00 2001
From: Paulus Schoutsen <balloob@gmail.com>
Date: Mon, 13 Apr 2020 17:41:01 -0700
Subject: [PATCH] Dump states in event handler for HA_Stop (#33974)

* Dump states in event handler for HA_Stop

* Fix type
---
 homeassistant/core.py                    | 4 +++-
 homeassistant/helpers/entity_platform.py | 4 +---
 homeassistant/helpers/event.py           | 6 ++++--
 homeassistant/helpers/restore_state.py   | 7 +++----
 4 files changed, 11 insertions(+), 10 deletions(-)

diff --git a/homeassistant/core.py b/homeassistant/core.py
index b3f4a68e37f..d35e7f48cc7 100644
--- a/homeassistant/core.py
+++ b/homeassistant/core.py
@@ -361,7 +361,9 @@ class HomeAssistant:
         self._track_task = False
 
     @callback
-    def async_run_job(self, target: Callable[..., None], *args: Any) -> None:
+    def async_run_job(
+        self, target: Callable[..., Union[None, Awaitable]], *args: Any
+    ) -> None:
         """Run a job from within the event loop.
 
         This method must be run in the event loop.
diff --git a/homeassistant/helpers/entity_platform.py b/homeassistant/helpers/entity_platform.py
index bf6db55a2da..e4d52aaa3a1 100644
--- a/homeassistant/helpers/entity_platform.py
+++ b/homeassistant/helpers/entity_platform.py
@@ -300,9 +300,7 @@ class EntityPlatform:
             return
 
         self._async_unsub_polling = async_track_time_interval(
-            self.hass,
-            self._update_entity_states,  # type: ignore
-            self.scan_interval,
+            self.hass, self._update_entity_states, self.scan_interval,
         )
 
     async def _async_add_entity(
diff --git a/homeassistant/helpers/event.py b/homeassistant/helpers/event.py
index 27ab35dbb3c..266cb150e0a 100644
--- a/homeassistant/helpers/event.py
+++ b/homeassistant/helpers/event.py
@@ -1,7 +1,7 @@
 """Helpers for listening to events."""
 from datetime import datetime, timedelta
 import functools as ft
-from typing import Any, Callable, Dict, Iterable, Optional, Union, cast
+from typing import Any, Awaitable, Callable, Dict, Iterable, Optional, Union, cast
 
 import attr
 
@@ -274,7 +274,9 @@ call_later = threaded_listener_factory(async_call_later)
 @callback
 @bind_hass
 def async_track_time_interval(
-    hass: HomeAssistant, action: Callable[..., None], interval: timedelta
+    hass: HomeAssistant,
+    action: Callable[..., Union[None, Awaitable]],
+    interval: timedelta,
 ) -> CALLBACK_TYPE:
     """Add a listener that fires repetitively at every timedelta interval."""
     remove = None
diff --git a/homeassistant/helpers/restore_state.py b/homeassistant/helpers/restore_state.py
index 57da6960078..c75d9c840ed 100644
--- a/homeassistant/helpers/restore_state.py
+++ b/homeassistant/helpers/restore_state.py
@@ -171,14 +171,13 @@ class RestoreStateData:
     def async_setup_dump(self, *args: Any) -> None:
         """Set up the restore state listeners."""
 
-        @callback
-        def _async_dump_states(*_: Any) -> None:
-            self.hass.async_create_task(self.async_dump_states())
+        async def _async_dump_states(*_: Any) -> None:
+            await self.async_dump_states()
 
         # Dump the initial states now. This helps minimize the risk of having
         # old states loaded by overwriting the last states once Home Assistant
         # has started and the old states have been read.
-        _async_dump_states()
+        self.hass.async_create_task(_async_dump_states())
 
         # Dump states periodically
         async_track_time_interval(self.hass, _async_dump_states, STATE_DUMP_INTERVAL)