Fix flaky history test

tests/components/test_history.py

@@ -5,11 +5,10 @@ tests.test_component_history
 Tests the history component.
 """
 # pylint: disable=protected-access,too-many-public-methods
-import time
+from datetime import timedelta
 import os
 import unittest
 from unittest.mock import patch
-from datetime import timedelta
 
 import homeassistant.core as ha
 import homeassistant.util.dt as dt_util
@@ -25,21 +24,24 @@ class TestComponentHistory(unittest.TestCase):
     def setUp(self):  # pylint: disable=invalid-name
         """ Init needed objects. """
         self.hass = get_test_home_assistant(1)
-        self.init_rec = False
 
     def tearDown(self):  # pylint: disable=invalid-name
         """ Stop down stuff we started. """
         self.hass.stop()
 
-        if self.init_rec:
-            recorder._INSTANCE.block_till_done()
-            os.remove(self.hass.config.path(recorder.DB_FILE))
+        db_path = self.hass.config.path(recorder.DB_FILE)
+        if os.path.isfile(db_path):
+            os.remove(db_path)
 
     def init_recorder(self):
         recorder.setup(self.hass, {})
         self.hass.start()
+        self.wait_recording_done()
+
+    def wait_recording_done(self):
+        """ Block till recording is done. """
+        self.hass.pool.block_till_done()
         recorder._INSTANCE.block_till_done()
-        self.init_rec = True
 
     def test_setup(self):
         """ Test setup method of history. """
@@ -56,12 +58,11 @@ class TestComponentHistory(unittest.TestCase):
         for i in range(7):
             self.hass.states.set(entity_id, "State {}".format(i))
 
+            self.wait_recording_done()
+
             if i > 1:
                 states.append(self.hass.states.get(entity_id))
 
-            self.hass.pool.block_till_done()
-            recorder._INSTANCE.block_till_done()
-
         self.assertEqual(
             list(reversed(states)), history.last_5_states(entity_id))
 
@@ -70,22 +71,9 @@ class TestComponentHistory(unittest.TestCase):
         self.init_recorder()
         states = []
 
-        for i in range(5):
-            state = ha.State(
-                'test.point_in_time_{}'.format(i % 5),
-                "State {}".format(i),
-                {'attribute_test': i})
-
-            mock_state_change_event(self.hass, state)
-            self.hass.pool.block_till_done()
-
-            states.append(state)
-
-        recorder._INSTANCE.block_till_done()
-
-        point = dt_util.utcnow() + timedelta(seconds=1)
-
-        with patch('homeassistant.util.dt.utcnow', return_value=point):
+        now = dt_util.utcnow()
+        with patch('homeassistant.components.recorder.dt_util.utcnow',
+                   return_value=now):
             for i in range(5):
                 state = ha.State(
                     'test.point_in_time_{}'.format(i % 5),
@@ -93,16 +81,32 @@ class TestComponentHistory(unittest.TestCase):
                     {'attribute_test': i})
 
                 mock_state_change_event(self.hass, state)
-                self.hass.pool.block_till_done()
+
+                states.append(state)
+
+            self.wait_recording_done()
+
+        future = now + timedelta(seconds=1)
+        with patch('homeassistant.components.recorder.dt_util.utcnow',
+                   return_value=future):
+            for i in range(5):
+                state = ha.State(
+                    'test.point_in_time_{}'.format(i % 5),
+                    "State {}".format(i),
+                    {'attribute_test': i})
+
+                mock_state_change_event(self.hass, state)
+
+            self.wait_recording_done()
 
         # Get states returns everything before POINT
         self.assertEqual(states,
-                         sorted(history.get_states(point),
+                         sorted(history.get_states(future),
                                 key=lambda state: state.entity_id))
 
         # Test get_state here because we have a DB setup
         self.assertEqual(
-            states[0], history.get_state(point, states[0].entity_id))
+            states[0], history.get_state(future, states[0].entity_id))
 
     def test_state_changes_during_period(self):
         self.init_recorder()
@@ -110,19 +114,20 @@ class TestComponentHistory(unittest.TestCase):
 
         def set_state(state):
             self.hass.states.set(entity_id, state)
-            self.hass.pool.block_till_done()
-            recorder._INSTANCE.block_till_done()
-
+            self.wait_recording_done()
             return self.hass.states.get(entity_id)
 
-        set_state('idle')
-        set_state('YouTube')
-
         start = dt_util.utcnow()
         point = start + timedelta(seconds=1)
         end = point + timedelta(seconds=1)
 
-        with patch('homeassistant.util.dt.utcnow', return_value=point):
+        with patch('homeassistant.components.recorder.dt_util.utcnow',
+                   return_value=start):
+            set_state('idle')
+            set_state('YouTube')
+
+        with patch('homeassistant.components.recorder.dt_util.utcnow',
+                   return_value=point):
             states = [
                 set_state('idle'),
                 set_state('Netflix'),
@@ -130,10 +135,11 @@ class TestComponentHistory(unittest.TestCase):
                 set_state('YouTube'),
             ]
 
-        with patch('homeassistant.util.dt.utcnow', return_value=end):
+        with patch('homeassistant.components.recorder.dt_util.utcnow',
+                   return_value=end):
             set_state('Netflix')
             set_state('Plex')
 
-        self.assertEqual(
-            {entity_id: states},
-            history.state_changes_during_period(start, end, entity_id))
+        hist = history.state_changes_during_period(start, end, entity_id)
+
+        self.assertEqual(states, hist[entity_id])