fixing indentation, adding tests to confirm connection behavior

xginn8 · xginn8 · commit 93b91e461114 · 2017-04-11T23:39:10.000-04:00
diff --git a/influxdb/client.py b/influxdb/client.py
@@ -252,6 +252,9 @@ def request(self, url, method='GET', params=None, data=None,
                 if self._retries != 0:
                     retry = _try < self._retries
 
+        else:
+            raise requests.exceptions.ConnectionError
+
         if 500 <= response.status_code < 600:
             raise InfluxDBServerError(response.content)
         elif response.status_code == expected_response_code:
diff --git a/influxdb/influxdb08/client.py b/influxdb/influxdb08/client.py
@@ -248,11 +248,13 @@ def request(self, url, method='GET', params=None, data=None,
                     timeout=self._timeout
                 )
                 break
-            except (requests.exceptions.ConnectionError,
-                    requests.exceptions.Timeout) as e:
+             except (requests.exceptions.ConnectionError,
+                     requests.exceptions.Timeout) as e:
                 _try += 1
                 if self._retries != 0:
                     retry = _try < self._retries
+        else:
+            raise requests.exceptions.ConnectionError
 
         if response.status_code == expected_response_code:
             return response
diff --git a/influxdb/tests/client_test.py b/influxdb/tests/client_test.py
@@ -645,6 +645,60 @@ def connection_error(self, *args, **kwargs):
         with self.assertRaises(requests.exceptions.ConnectionError):
             cli.write_points(self.dummy_points)
 
+    @mock.patch('requests.Session.request')
+    def test_random_request_retry(self, mock_request):
+        """Tests that a random number of connection errors will be handled"""
+
+        class CustomMock(object):
+            def __init__(self, retries):
+                self.i = 0
+                self.retries = retries
+
+            def connection_error(self, *args, **kwargs):
+                self.i += 1
+
+                if self.i < self.retries:
+                    raise requests.exceptions.ConnectionError
+                else:
+                    r = requests.Response()
+                    r.status_code = 204
+                    return r
+
+        retries = random.randint(1, 100)
+        mock_request.side_effect = CustomMock(retries).connection_error
+
+        cli = InfluxDBClient(database='db', retries=retries)
+        cli.write_points(
+            self.dummy_points
+        )
+
+    @mock.patch('requests.Session.request')
+    def test_random_request_retry_raises(self, mock_request):
+        """Tests that a random number of connection errors plus one will be not handled"""
+
+        class CustomMock(object):
+            def __init__(self, retries):
+                self.i = 0
+                self.retries = retries
+
+            def connection_error(self, *args, **kwargs):
+                self.i += 1
+
+                if self.i < self.retries + 1:
+                    raise requests.exceptions.ConnectionError
+                else:
+                    r = requests.Response()
+                    r.status_code = 200
+                    return r
+
+        retries = random.randint(1, 100)
+        mock_request.side_effect = CustomMock(retries).connection_error
+
+        cli = InfluxDBClient(database='db', retries=retries)
+
+        with self.assertRaises(requests.exceptions.ConnectionError):
+            cli.write_points(self.dummy_points)
+
     def test_get_list_users(self):
         example_response = (
             '{"results":[{"series":[{"columns":["user","admin"],'
