Merge pull request #39 from naved001/handle-network-failures

Handle network failures when gathering metrics.

openshift_metrics/tests/test_utils.py

@@ -10,7 +10,7 @@
 #   License for the specific language governing permissions and limitations
 #   under the License.
 #
-
+from requests.exceptions import ConnectionError
 import tempfile
 from unittest import TestCase, mock
 
@@ -19,7 +19,7 @@
 
 class TestQueryMetric(TestCase):
 
-    @mock.patch('requests.get')
+    @mock.patch('requests.Session.get')
     @mock.patch('time.sleep')
     def test_query_metric(self, mock_sleep, mock_get):
         mock_response = mock.Mock(status_code=200)
@@ -28,19 +28,28 @@ def test_query_metric(self, mock_sleep, mock_get):
         }}
         mock_get.return_value = mock_response
 
-        metrics = utils.query_metric('fake-url', 'fake-token', 'fake-metric', '2022-03-14', '2022-03-14')
+        metrics = utils.query_metric('https://fake-url', 'fake-token', 'fake-metric', '2022-03-14', '2022-03-14')
         self.assertEqual(metrics, "this is data")
         self.assertEqual(mock_get.call_count, 1)
 
-    @mock.patch('requests.get')
+    @mock.patch('requests.Session.get')
     @mock.patch('time.sleep')
     def test_query_metric_exception(self, mock_sleep, mock_get):
         mock_get.return_value = mock.Mock(status_code=404)
 
-        self.assertRaises(Exception, utils.query_metric, 'fake-url', 'fake-token',
+        self.assertRaises(Exception, utils.query_metric, 'https://fake-url', 'fake-token',
                           'fake-metric', '2022-03-14', '2022-03-14')
         self.assertEqual(mock_get.call_count, 3)
 
+    @mock.patch('requests.Session.get')
+    @mock.patch('time.sleep')
+    def test_query_metric_connection_error(self, mock_sleep, mock_get):
+        mock_get.side_effect = [ConnectionError]
+        self.assertRaises(ConnectionError, utils.query_metric, 'https://fake-url', 'fake-token',
+                  'fake-metric', '2022-03-14', '2022-03-14')
+        self.assertEqual(mock_get.call_count, 1)
+
+
 class TestGetNamespaceAnnotations(TestCase):
 
     @mock.patch('openshift_metrics.utils.requests.post')

openshift_metrics/utils.py

@@ -21,6 +21,8 @@
 import requests
 import boto3
 
+from urllib3.util.retry import Retry
+from requests.adapters import HTTPAdapter
 
 # GPU types
 GPU_A100 = "nvidia.com/gpu_A100"
@@ -104,10 +106,17 @@ def query_metric(openshift_url, token, metric, report_start_date, report_end_dat
     data = None
     headers = {"Authorization": f"Bearer {token}"}
     day_url_vars = f"start={report_start_date}T00:00:00Z&end={report_end_date}T23:59:59Z"
+    url = f"{openshift_url}/api/v1/query_range?query={metric}&{day_url_vars}&step={STEP_MIN}m"
+
+    retries = Retry(total=3, backoff_factor=1, status_forcelist=[429, 500, 502, 503, 504])
+    session = requests.Session()
+    session.mount("https://", HTTPAdapter(max_retries=retries))
+
     print(f"Retrieving metric: {metric}")
+
     for _ in range(3):
-        url = f"{openshift_url}/api/v1/query_range?query={metric}&{day_url_vars}&step={STEP_MIN}m"
-        response = requests.get(url, headers=headers, verify=True)
+        response = session.get(url, headers=headers, verify=True)
+
         if response.status_code != 200:
             print(f"{response.status_code} Response: {response.reason}")
         else:
@@ -116,6 +125,7 @@ def query_metric(openshift_url, token, metric, report_start_date, report_end_dat
                 break
             print("Empty result set")
         time.sleep(3)
+
     if not data:
         raise EmptyResultError(f"Error retrieving metric: {metric}")
     return data