1 files changed, 123 insertions, 0 deletions
diff --git a/tests/utils/test_rest_util.py b/tests/utils/test_rest_util.py
index 52674da..bde0c3d 100644
--- a/tests/utils/test_rest_util.py
+++ b/tests/utils/test_rest_util.py
@@ -24,9 +24,70 @@ from azurelinuxagent.common.exception import HttpError, \
 import azurelinuxagent.common.utils.restutil as restutil
 
 from azurelinuxagent.common.future import httpclient, ustr
+
 from tests.tools import *
 
 
+class TestIOErrorCounter(AgentTestCase):
+    def test_increment_hostplugin(self):
+        restutil.IOErrorCounter.reset()
+        restutil.IOErrorCounter.set_protocol_endpoint()
+
+        restutil.IOErrorCounter.increment(
+            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+
+        counts = restutil.IOErrorCounter.get_and_reset()
+        self.assertEqual(1, counts["hostplugin"])
+        self.assertEqual(0, counts["protocol"])
+        self.assertEqual(0, counts["other"])
+
+    def test_increment_protocol(self):
+        restutil.IOErrorCounter.reset()
+        restutil.IOErrorCounter.set_protocol_endpoint()
+
+        restutil.IOErrorCounter.increment(
+            restutil.DEFAULT_PROTOCOL_ENDPOINT, 80)
+
+        counts = restutil.IOErrorCounter.get_and_reset()
+        self.assertEqual(0, counts["hostplugin"])
+        self.assertEqual(1, counts["protocol"])
+        self.assertEqual(0, counts["other"])
+
+    def test_increment_other(self):
+        restutil.IOErrorCounter.reset()
+        restutil.IOErrorCounter.set_protocol_endpoint()
+
+        restutil.IOErrorCounter.increment(
+            '169.254.169.254', 80)
+
+        counts = restutil.IOErrorCounter.get_and_reset()
+        self.assertEqual(0, counts["hostplugin"])
+        self.assertEqual(0, counts["protocol"])
+        self.assertEqual(1, counts["other"])
+
+    def test_get_and_reset(self):
+        restutil.IOErrorCounter.reset()
+        restutil.IOErrorCounter.set_protocol_endpoint()
+
+        restutil.IOErrorCounter.increment(
+            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+        restutil.IOErrorCounter.increment(
+            restutil.DEFAULT_PROTOCOL_ENDPOINT, restutil.HOST_PLUGIN_PORT)
+        restutil.IOErrorCounter.increment(
+            restutil.DEFAULT_PROTOCOL_ENDPOINT, 80)
+        restutil.IOErrorCounter.increment(
+            '169.254.169.254', 80)
+        restutil.IOErrorCounter.increment(
+            '169.254.169.254', 80)
+
+        counts = restutil.IOErrorCounter.get_and_reset()
+        self.assertEqual(2, counts.get("hostplugin"))
+        self.assertEqual(1, counts.get("protocol"))
+        self.assertEqual(2, counts.get("other"))
+        self.assertEqual(
+           {"hostplugin":0, "protocol":0, "other":0},
+            restutil.IOErrorCounter._counts)
+
 class TestHttpOperations(AgentTestCase):
     def test_parse_url(self):
         test_uri = "http://abc.def/ghi#hash?jkl=mn"
@@ -268,6 +329,68 @@ class TestHttpOperations(AgentTestCase):
 
     @patch("time.sleep")
     @patch("azurelinuxagent.common.utils.restutil._http_request")
+    def test_http_request_retries_with_fibonacci_delay(self, _http_request, _sleep):
+        # Ensure the code is not a throttle code
+        self.assertFalse(httpclient.BAD_GATEWAY in restutil.THROTTLE_CODES)
+
+        _http_request.side_effect = [
+                Mock(status=httpclient.BAD_GATEWAY)
+                    for i in range(restutil.DEFAULT_RETRIES)
+            ] + [Mock(status=httpclient.OK)]
+
+        restutil.http_get("https://foo.bar",
+                            max_retry=restutil.DEFAULT_RETRIES+1)
+
+        self.assertEqual(restutil.DEFAULT_RETRIES+1, _http_request.call_count)
+        self.assertEqual(restutil.DEFAULT_RETRIES, _sleep.call_count)
+        self.assertEqual(
+            [
+                call(restutil._compute_delay(i+1, restutil.DELAY_IN_SECONDS))
+                    for i in range(restutil.DEFAULT_RETRIES)],
+            _sleep.call_args_list)
+
+    @patch("time.sleep")
+    @patch("azurelinuxagent.common.utils.restutil._http_request")
+    def test_http_request_retries_with_constant_delay_when_throttled(self, _http_request, _sleep):
+        # Ensure the code is a throttle code
+        self.assertTrue(httpclient.SERVICE_UNAVAILABLE in restutil.THROTTLE_CODES)
+
+        _http_request.side_effect = [
+                Mock(status=httpclient.SERVICE_UNAVAILABLE)
+                    for i in range(restutil.DEFAULT_RETRIES)
+            ] + [Mock(status=httpclient.OK)]
+
+        restutil.http_get("https://foo.bar",
+                            max_retry=restutil.DEFAULT_RETRIES+1)
+
+        self.assertEqual(restutil.DEFAULT_RETRIES+1, _http_request.call_count)
+        self.assertEqual(restutil.DEFAULT_RETRIES, _sleep.call_count)
+        self.assertEqual(
+            [call(1) for i in range(restutil.DEFAULT_RETRIES)],
+            _sleep.call_args_list)
+
+    @patch("time.sleep")
+    @patch("azurelinuxagent.common.utils.restutil._http_request")
+    def test_http_request_retries_for_safe_minimum_number_when_throttled(self, _http_request, _sleep):
+        # Ensure the code is a throttle code
+        self.assertTrue(httpclient.SERVICE_UNAVAILABLE in restutil.THROTTLE_CODES)
+
+        _http_request.side_effect = [
+                Mock(status=httpclient.SERVICE_UNAVAILABLE)
+                    for i in range(restutil.THROTTLE_RETRIES-1)
+            ] + [Mock(status=httpclient.OK)]
+
+        restutil.http_get("https://foo.bar",
+                            max_retry=1)
+
+        self.assertEqual(restutil.THROTTLE_RETRIES, _http_request.call_count)
+        self.assertEqual(restutil.THROTTLE_RETRIES-1, _sleep.call_count)
+        self.assertEqual(
+            [call(1) for i in range(restutil.THROTTLE_RETRIES-1)],
+            _sleep.call_args_list)
+
+    @patch("time.sleep")
+    @patch("azurelinuxagent.common.utils.restutil._http_request")
     def test_http_request_raises_for_bad_request(self, _http_request, _sleep):
         _http_request.side_effect = [
             Mock(status=httpclient.BAD_REQUEST)