Merge pull request #557 from cmu-delphi/retries

Add retry to Epidata request

Author: krivard

integrations/client/test_delphi_epidata.py

@@ -3,6 +3,7 @@
 # standard library
 import unittest
 from unittest.mock import patch, MagicMock
+from json import JSONDecodeError
 
 # third party
 from aiohttp.client_exceptions import ClientResponseError
@@ -293,13 +294,38 @@ def test_request_method(self, get, post):
       get.assert_called_once()
       post.assert_not_called()
     with self.subTest(name='post request'):
+      get.reset_mock()
       mock_response = MagicMock()
       mock_response.status_code = 414
       get.return_value = mock_response
       Epidata.covidcast('src', 'sig', 'day', 'county', 20200414, '01234')
-      self.assertEqual(get.call_count, 2)  # one from post test and one from get test
+      get.assert_called_once()
       post.assert_called_once()
 
+  @patch('requests.get')
+  def test_retry_request(self, get):
+    """Test that a GET request is default and POST is used if a 414 is returned."""
+    with self.subTest(name='test successful retry'):
+      mock_response = MagicMock()
+      mock_response.status_code = 200
+      get.side_effect = [JSONDecodeError('Expecting value', "",  0), mock_response]
+      response = Epidata._request(None)
+      self.assertEqual(get.call_count, 2)
+      self.assertEqual(response, mock_response.json())
+
+    with self.subTest(name='test retry'):
+      get.reset_mock()
+      mock_response = MagicMock()
+      mock_response.status_code = 200
+      get.side_effect = [JSONDecodeError('Expecting value', "",  0),
+                         JSONDecodeError('Expecting value', "",  0),
+                         mock_response]
+      response = Epidata._request(None)
+      self.assertEqual(get.call_count, 2)  # 2 from previous test + 2 from this one
+      self.assertEqual(response,
+                       {'result': 0, 'message': 'error: Expecting value: line 1 column 1 (char 0)'}
+                       )
+
   def test_geo_value(self):
     """test different variants of geo types: single, *, multi."""
 

requirements.dev.txt

@@ -3,4 +3,5 @@ black>=20.8b1
 sqlalchemy-stubs>=0.3
 mypy>=0.790
 pytest
+tenacity==7.0.0
 bump2version

requirements.txt

@@ -5,4 +5,5 @@ python-dotenv==0.15.0
 orjson==3.4.7
 pandas==1.2.3
 scipy==1.6.2
+tenacity==7.0.0
 newrelic

src/client/delphi_epidata.py

@@ -11,7 +11,7 @@
 # External modules
 import requests
 import asyncio
-import warnings
+from tenacity import retry, stop_after_attempt
 
 from aiohttp import ClientSession, TCPConnector
 from pkg_resources import get_distribution, DistributionNotFound
@@ -52,7 +52,15 @@ def _list(values):
       values = [values]
     return ','.join([Epidata._listitem(value) for value in values])
 
-  # Helper function to request and parse epidata
+  @staticmethod
+  @retry(reraise=True, stop=stop_after_attempt(2))
+  def _request_with_retry(params):
+    """Make request with a retry if an exception is thrown."""
+    req = requests.get(Epidata.BASE_URL, params, headers=_HEADERS)
+    if req.status_code == 414:
+      req = requests.post(Epidata.BASE_URL, params, headers=_HEADERS)
+    return req
+
   @staticmethod
   def _request(params):
     """Request and parse epidata.
@@ -62,13 +70,8 @@ def _request(params):
     long and returns a 414.
     """
     try:
-      # API call
-      req = requests.get(Epidata.BASE_URL, params, headers=_HEADERS)
-      if req.status_code == 414:
-        req = requests.post(Epidata.BASE_URL, params, headers=_HEADERS)
-      return req.json()
+      return Epidata._request_with_retry(params).json()
     except Exception as e:
-      # Something broke
       return {'result': 0, 'message': 'error: ' + str(e)}
 
   # Raise an Exception on error, otherwise return epidata

src/client/packaging/pypi/setup.py

@@ -16,6 +16,7 @@
   install_requires=[
     'aiohttp',
     'requests>=2.7.0',
+    'tenacity'
   ],
   classifiers=[
     'Programming Language :: Python',