Add unittests for RetryOnError

Signed-off-by: Jan Losinski <losinski@wh2.tu-dresden.de>

tests/util/test_init.py

@@ -1,11 +1,14 @@
 """Test Home Assistant util methods."""
 import unittest
-from unittest.mock import patch
+from unittest.mock import patch, MagicMock
 from datetime import datetime, timedelta
 
+from homeassistant import core as ha
 from homeassistant import util
 import homeassistant.util.dt as dt_util
 
+from tests.common import get_test_home_assistant
+
 
 class TestUtil(unittest.TestCase):
     """Test util methods."""
@@ -266,3 +269,78 @@ def goodbye(self):
 
         self.assertTrue(tester.hello())
         self.assertTrue(tester.goodbye())
+
+
+class TestRetryOnErrorDecoratoru(unittest.TestCase):
+    """Test the RetryOnError decorator."""
+
+    def setUp(self):
+        """Setup things to be run when tests are started."""
+        self.hass = get_test_home_assistant()
+
+    def tearDown(self):
+        """Clear data."""
+        self.hass.stop()
+
+    def test_no_retry(self):
+        """Test that it does not retry if configured."""
+        mock_method = MagicMock()
+        wrapped = util.RetryOnError(self.hass)(mock_method)
+        wrapped(1, 2, test=3)
+        self.assertEqual(mock_method.call_count, 1)
+        mock_method.assert_called_with(1, 2, test=3)
+
+        mock_method.side_effect = Exception()
+        self.assertRaises(Exception, wrapped, 1, 2, test=3)
+        self.assertEqual(mock_method.call_count, 2)
+        mock_method.assert_called_with(1, 2, test=3)
+
+    def test_single_retry(self):
+        """Test that retry stops after a single try if configured."""
+        mock_method = MagicMock()
+        retryer = util.RetryOnError(self.hass, retry_limit=1)
+        wrapped = retryer(mock_method)
+        wrapped(1, 2, test=3)
+        self.assertEqual(mock_method.call_count, 1)
+        mock_method.assert_called_with(1, 2, test=3)
+
+        start = dt_util.utcnow()
+        shifted_time = start + (timedelta(seconds=20 + 1))
+        self.hass.bus.fire(ha.EVENT_TIME_CHANGED,
+                           {ha.ATTR_NOW: shifted_time})
+        self.hass.block_till_done()
+        self.assertEqual(mock_method.call_count, 1)
+
+        mock_method.side_effect = Exception()
+        wrapped(1, 2, test=3)
+        self.assertEqual(mock_method.call_count, 2)
+        mock_method.assert_called_with(1, 2, test=3)
+
+        for cnt in range(3):
+            start = dt_util.utcnow()
+            shifted_time = start + (timedelta(seconds=20 + 1))
+            self.hass.bus.fire(ha.EVENT_TIME_CHANGED,
+                               {ha.ATTR_NOW: shifted_time})
+            self.hass.block_till_done()
+            self.assertEqual(mock_method.call_count, 3)
+            mock_method.assert_called_with(1, 2, test=3)
+
+    def test_multi_retry(self):
+        """Test that multiple retries work."""
+        mock_method = MagicMock()
+        retryer = util.RetryOnError(self.hass, retry_limit=4)
+        wrapped = retryer(mock_method)
+        mock_method.side_effect = Exception()
+
+        wrapped(1, 2, test=3)
+        self.assertEqual(mock_method.call_count, 1)
+        mock_method.assert_called_with(1, 2, test=3)
+
+        for cnt in range(3):
+            start = dt_util.utcnow()
+            shifted_time = start + (timedelta(seconds=20 + 1))
+            self.hass.bus.fire(ha.EVENT_TIME_CHANGED,
+                               {ha.ATTR_NOW: shifted_time})
+            self.hass.block_till_done()
+            self.assertEqual(mock_method.call_count, cnt + 2)
+            mock_method.assert_called_with(1, 2, test=3)