Merge pull request #145 from jcoatgoogle/jco/batch_callback

Add batch request callback hook
google · Feb 28, 2017 · 1f57d98 · 1f57d98
2 parents b71c414 + 4e88a3d
commit 1f57d98
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/apitools/base/py/batch.py b/apitools/base/py/batch.py
@@ -182,7 +182,7 @@ def Add(self, service, method, request, global_params=None):
         self.api_requests.append(api_request)
 
     def Execute(self, http, sleep_between_polls=5, max_retries=5,
-                max_batch_size=None):
+                max_batch_size=None, batch_request_callback=None):
         """Execute all of the requests in the batch.
 
         Args:
@@ -194,6 +194,8 @@ def Execute(self, http, sleep_between_polls=5, max_retries=5,
               exception, whatever it happened to be.
           max_batch_size: int, if specified requests will be split in batches
               of given size.
+          batch_request_callback: function of (http_response, exception) passed
+              to BatchHttpRequest which will be run on any given results.
 
         Returns:
           List of ApiCalls.
@@ -209,7 +211,10 @@ def Execute(self, http, sleep_between_polls=5, max_retries=5,
             for i in range(0, len(requests), batch_size):
                 # Create a batch_http_request object and populate it with
                 # incomplete requests.
-                batch_http_request = BatchHttpRequest(batch_url=self.batch_url)
+                batch_http_request = BatchHttpRequest(
+                    batch_url=self.batch_url,
+                    callback=batch_request_callback
+                )
                 for request in itertools.islice(requests,
                                                 i, i + batch_size):
                     batch_http_request.Add(

diff --git a/apitools/base/py/batch_test.py b/apitools/base/py/batch_test.py
@@ -150,9 +150,18 @@ def testRequestServiceUnavailable(self):
                                   exceptions.HttpError)
 
     def testSingleRequestInBatch(self):
+        desired_url = 'https://www.example.com'
+
+        callback_was_called = []
+        def _Callback(response, exception):
+            self.assertEqual({'status': '200'}, response.info)
+            self.assertEqual('content', response.content)
+            self.assertEqual(desired_url, response.request_url)
+            self.assertIsNone(exception)
+            callback_was_called.append(1)
+
         mock_service = FakeService()
 
-        desired_url = 'https://www.example.com'
         batch_api_request = batch.BatchApiRequest(batch_url=desired_url)
         # The request to be added. The actual request sent will be somewhat
         # larger, as this is added to a batch.
@@ -185,7 +194,8 @@ def testSingleRequestInBatch(self):
                 'desired_request': desired_request,
             })
 
-            api_request_responses = batch_api_request.Execute(FakeHttp())
+            api_request_responses = batch_api_request.Execute(
+                FakeHttp(), batch_request_callback=_Callback)
 
             self.assertEqual(1, len(api_request_responses))
             self.assertEqual(1, mock_request.call_count)
@@ -196,6 +206,7 @@ def testSingleRequestInBatch(self):
             self.assertEqual({'status': '200'}, response.info)
             self.assertEqual('content', response.content)
             self.assertEqual(desired_url, response.request_url)
+        self.assertEquals(1, len(callback_was_called))
 
     def _MakeResponse(self, number_of_parts):
         return http_wrapper.Response(