Merge pull request #30 from Yelp/update-batch-url-behavior

Update BatchRequest batch_url to get_job_group_url
Yelp · May 17, 2021 · d6aae30 · d6aae30
2 parents 8cb921e + 6c89bac
commit d6aae30
Show file tree

Hide file tree

Showing 5 changed files with 110 additions and 15 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -0,0 +1,78 @@
+# Changelog
+All notable changes to this project will be documented in this file.
+
+The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
+and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
+
+## [Unreleased]
+
+## [8.0.0] - 2021-04-16
+
+### Changed
+
+`BatchRequest` now accepts a function, `get_job_group_url`, in place of `batch_url`. This allows for increased flexibility in routing hypernova requests.
+
+You'll need to update your `pyramid_hypernova.get_batch_url` to the `pyramid_hypernova.get_job_group_url` setting.
+
+Before:
+```py
+def get_batch_url():
+    return 'https://localhost:8080/batch'
+
+registry.settings.update({
+    'pyramid_hypernova.get_batch_url': get_batch_url,
+})
+```
+
+After:
+```py
+def get_job_group_url():
+    return 'https://localhost:8080/batch'
+
+registry.settings.update({
+    'pyramid_hypernova.get_job_group_url': get_job_group_url,
+})
+```
+
+-----
+
+Additionally, If you were manually supplying a BatchRequest factory via the `pyramid_hypernova.batch_request_factory` setting, you'll need
+to update its API to accept `get_job_group_url` and pass it along to the `BatchRequest` class. Example:
+
+Before:
+```py
+def batch_request_factory(batch_url, plugin_controller, json_encoder, pyramid_request):
+    return BatchRequest(
+        batch_url=batch_url,
+        plugin_controller=plugin_controller,
+        json_encoder=json_encoder,
+        pyramid_request=pyramid_request,
+    )
+
+registry.settings.update({
+    'pyramid_hypernova.batch_request_factory': batch_request_factory,
+})
+```
+
+After:
+```py
+def batch_request_factory(get_job_group_url, plugin_controller, json_encoder, pyramid_request):
+    return BatchRequest(
+        get_job_group_url=get_job_group_url,
+        plugin_controller=plugin_controller,
+        json_encoder=json_encoder,
+        pyramid_request=pyramid_request,
+    )
+
+registry.settings.update({
+    'pyramid_hypernova.batch_request_factory': batch_request_factory,
+})
+```
+
+`get_job_group_url` will be supplied two parameters - the **job group** that pyramid_hypernova will be making a request for and the **pyramid request**:
+
+```py
+def get_batch_url(job_group, pyramid_request): ...
+```
+
+This will be called right before send _on every job group request_. If you're doing an expensive operation to retrieve this url, consider memoizing it.
diff --git a/pyramid_hypernova/batch.py b/pyramid_hypernova/batch.py
@@ -50,8 +50,15 @@ def create_job_groups(jobs, max_batch_size):
 
 class BatchRequest(object):
 
-    def __init__(self, batch_url, plugin_controller, pyramid_request, max_batch_size=None, json_encoder=JSONEncoder()):
-        self.batch_url = batch_url
+    def __init__(
+        self,
+        get_job_group_url,
+        plugin_controller,
+        pyramid_request,
+        max_batch_size=None,
+        json_encoder=JSONEncoder()
+    ):
+        self.get_job_group_url = get_job_group_url
         self.jobs = {}
         self.plugin_controller = plugin_controller
         self.max_batch_size = max_batch_size
@@ -156,8 +163,9 @@ def submit(self):
             synchronous = len(job_groups) == 1
 
             for job_group in job_groups:
+                batch_url = self.get_job_group_url(job_group, self.pyramid_request)
                 request_headers = self.plugin_controller.transform_request_headers({}, self.pyramid_request)
-                query = HypernovaQuery(job_group, self.batch_url, self.json_encoder, synchronous, request_headers)
+                query = HypernovaQuery(job_group, batch_url, self.json_encoder, synchronous, request_headers)
                 query.send()
                 queries.append((job_group, query))
 

diff --git a/pyramid_hypernova/tweens.py b/pyramid_hypernova/tweens.py
@@ -38,7 +38,7 @@ def hypernova_tween(request):
 
 
 def configure_hypernova_batch(registry, request):
-    get_batch_url = registry.settings['pyramid_hypernova.get_batch_url']
+    get_job_group_url = registry.settings['pyramid_hypernova.get_job_group_url']
 
     plugins = registry.settings.get('pyramid_hypernova.plugins', [])
     plugin_controller = PluginController(plugins)
@@ -51,7 +51,7 @@ def configure_hypernova_batch(registry, request):
     json_encoder = registry.settings.get('pyramid_hypernova.json_encoder', JSONEncoder())
 
     return batch_request_factory(
-        batch_url=get_batch_url(),
+        get_job_group_url=get_job_group_url,
         plugin_controller=plugin_controller,
         json_encoder=json_encoder,
         pyramid_request=request,

diff --git a/tests/batch_test.py b/tests/batch_test.py
@@ -97,6 +97,11 @@ def spy_plugin_controller():
     return mock.Mock(wraps=plugin_controller)
 
 
+@pytest.fixture
+def spy_get_job_group_url():
+    return mock.Mock(return_value='http://localhost:8888')
+
+
 @pytest.fixture(params=[
     # (data, use_complex_json_encoder)
     ([{'key-1': 'value-1'}, {'key-2': 'value-2'}, {'key-3': 'value-3'}], False),
@@ -108,11 +113,11 @@ def test_data(request):
 
 
 @pytest.fixture(params=[None, 1, 2])
-def batch_request(spy_plugin_controller, test_data, request):
+def batch_request(spy_get_job_group_url, spy_plugin_controller, test_data, request):
     json_encoder = ComplexJSONEncoder() if test_data[1] else JSONEncoder()
     return BatchRequest(
-        'http://localhost:8888',
-        spy_plugin_controller,
+        get_job_group_url=spy_get_job_group_url,
+        plugin_controller=spy_plugin_controller,
         pyramid_request=pyramid.request.Request.blank('/'),
         max_batch_size=request.param,
         json_encoder=json_encoder,
@@ -127,7 +132,7 @@ def mock_hypernova_query():
 
 class TestBatchRequest(object):
 
-    def test_successful_batch_request(self, spy_plugin_controller, test_data, batch_request, mock_hypernova_query):
+    def test_successful_batch_request(self, spy_get_job_group_url, test_data, batch_request, mock_hypernova_query):
         data = test_data[0]
         token_1 = batch_request.render('component-1.js', data[0])
         token_2 = batch_request.render('component-2.js', data[1])
@@ -172,14 +177,18 @@ def test_successful_batch_request(self, spy_plugin_controller, test_data, batch_
         response = batch_request.submit()
 
         if batch_request.max_batch_size is None:
+            assert spy_get_job_group_url.call_count == 1
+            # get_job_group_url is supplied with the job group
+            assert len(spy_get_job_group_url.mock_calls[0].args[0]) == 3
             assert mock_hypernova_query.call_count == 1
         else:
             # Division (rounded-up) up to get total number of calls
             jobs_count = len(batch_request.jobs)
             max_batch_size = batch_request.max_batch_size
             batch_count = (jobs_count + (max_batch_size - 1)) // max_batch_size
+            assert spy_get_job_group_url.call_count == batch_count
             assert mock_hypernova_query.call_count == batch_count
-            mock_hypernova_query.assert_called_with(mock.ANY, mock.ANY, mock.ANY, batch_count == 1, {})
+            mock_hypernova_query.assert_called_with(mock.ANY, 'http://localhost:8888', mock.ANY, batch_count == 1, {})
 
         assert response == {
             token_1.identifier: JobResult(

diff --git a/tests/tweens_test.py b/tests/tweens_test.py
@@ -21,7 +21,7 @@ def mock_setup(self):
             text=str(self.token)
         )
 
-        mock_get_batch_url = mock.Mock(return_value='http://localhost:8888/batch')
+        self.mock_get_job_group_url = mock.Mock(return_value='http://localhost:8888/batch')
 
         self.mock_json_encoder = mock.Mock()
 
@@ -36,7 +36,7 @@ def mock_setup(self):
 
         self.mock_registry = mock.Mock()
         self.mock_registry.settings = {
-            'pyramid_hypernova.get_batch_url': mock_get_batch_url,
+            'pyramid_hypernova.get_job_group_url': self.mock_get_job_group_url,
             'pyramid_hypernova.batch_request_factory': self.mock_batch_request_factory,
             'pyramid_hypernova.json_encoder': self.mock_json_encoder,
         }
@@ -51,7 +51,7 @@ def test_tween_replaces_tokens_when_disable_hypernova_tween_not_set(self):
         response = self.tween(self.mock_request)
 
         self.mock_batch_request_factory.assert_called_once_with(
-            batch_url='http://localhost:8888/batch',
+            get_job_group_url=self.mock_get_job_group_url,
             plugin_controller=mock.ANY,
             json_encoder=self.mock_json_encoder,
             pyramid_request=self.mock_request,
@@ -65,7 +65,7 @@ def test_tween_replaces_tokens_when_disable_hypernova_tween_set_false(self):
         response = self.tween(self.mock_request)
 
         self.mock_batch_request_factory.assert_called_once_with(
-            batch_url='http://localhost:8888/batch',
+            get_job_group_url=self.mock_get_job_group_url,
             plugin_controller=mock.ANY,
             json_encoder=self.mock_json_encoder,
             pyramid_request=self.mock_request,
@@ -79,7 +79,7 @@ def test_tween_replaces_tokens_when_disable_hypernova_tween_set_true(self):
         response = self.tween(self.mock_request)
 
         self.mock_batch_request_factory.assert_called_once_with(
-            batch_url='http://localhost:8888/batch',
+            get_job_group_url=self.mock_get_job_group_url,
             plugin_controller=mock.ANY,
             json_encoder=self.mock_json_encoder,
             pyramid_request=self.mock_request,