Merge pull request #432 from cmu-delphi/async-limit

Have async raise 4xx/5xx status codes and add connection limit

Author: krivard

integrations/client/test_delphi_epidata.py

@@ -5,7 +5,9 @@
 from unittest.mock import patch, MagicMock
 
 # third party
+from aiohttp.client_exceptions import ClientResponseError
 import mysql.connector
+import pytest
 
 # first party
 from delphi.epidata.client.delphi_epidata import Epidata
@@ -15,6 +17,14 @@
 # py3tester coverage target
 __test_target__ = 'delphi.epidata.client.delphi_epidata'
 
+def fake_epidata_endpoint(func):
+  """This can be used as a decorator to enable a bogus Epidata endpoint to return 404 responses."""
+  def wrapper(*args):
+    Epidata.BASE_URL = 'http://delphi_web_epidata/epidata/fake_api.php'
+    func(*args)
+    Epidata.BASE_URL = 'http://delphi_web_epidata/epidata/api.php'
+  return wrapper
+
 
 class DelphiEpidataPythonClientTests(unittest.TestCase):
   """Tests the Python client."""
@@ -509,10 +519,25 @@ def test_async_epidata(self):
         'geo_value': '00000',
         'time_values': '20200414'
       }
-    ], batch_size=10)
-    responses = [i[0] for i in test_output]*12
+    ]*12, batch_size=10)
+    responses = [i[0] for i in test_output]
     # check response is same as standard covidcast call, using 24 calls to test batch sizing
     self.assertEqual(responses,
                      [Epidata.covidcast('src', 'sig', 'day', 'county', 20200414, '11111'),
                       Epidata.covidcast('src', 'sig', 'day', 'county', 20200414, '00000')]*12
                      )
+
+  @fake_epidata_endpoint
+  def test_async_epidata_fail(self):
+    with pytest.raises(ClientResponseError, match="404, message='Not Found'"):
+      Epidata.async_epidata([
+        {
+          'source': 'covidcast',
+          'data_source': 'src',
+          'signals': 'sig',
+          'time_type': 'day',
+          'geo_type': 'county',
+          'geo_value': '11111',
+          'time_values': '20200414'
+        }
+      ])

src/client/delphi_epidata.py

@@ -13,7 +13,7 @@
 import asyncio
 import warnings
 
-from aiohttp import ClientSession
+from aiohttp import ClientSession, TCPConnector
 from pkg_resources import get_distribution, DistributionNotFound
 
 # Obtain package version for the user-agent. Uses the installed version by
@@ -709,27 +709,26 @@ def covidcast_nowcast(
     return Epidata._request(params)
 
   @staticmethod
-  def async_epidata(param_list, batch_size=100):
+  def async_epidata(param_list, batch_size=50):
     """Make asynchronous Epidata calls for a list of parameters."""
     async def async_get(params, session):
       """Helper function to make Epidata GET requests."""
       async with session.get(Epidata.BASE_URL, params=params) as response:
+        response.raise_for_status()
         return await response.json(), params
 
     async def async_make_calls(param_combos):
       """Helper function to asynchronously make and aggregate Epidata GET requests."""
       tasks = []
-      async with ClientSession() as session:
+      connector = TCPConnector(limit=batch_size)
+      async with ClientSession(connector=connector) as session:
         for param in param_combos:
           task = asyncio.ensure_future(async_get(param, session))
           tasks.append(task)
         responses = await asyncio.gather(*tasks)
         return responses
 
-    batches = [param_list[i:i+batch_size] for i in range(0, len(param_list), batch_size)]
-    responses = []
-    for batch in batches:
-      loop = asyncio.get_event_loop()
-      future = asyncio.ensure_future(async_make_calls(batch))
-      responses += loop.run_until_complete(future)
+    loop = asyncio.get_event_loop()
+    future = asyncio.ensure_future(async_make_calls(param_list))
+    responses = loop.run_until_complete(future)
     return responses