diff --git a/pyramid_hypernova/tweens.py b/pyramid_hypernova/tweens.py index 9495a7d..ba094b1 100644 --- a/pyramid_hypernova/tweens.py +++ b/pyramid_hypernova/tweens.py @@ -8,26 +8,34 @@ def hypernova_tween_factory(handler, registry): registry = registry - def hypernova_tween(request): - request.hypernova_batch = configure_hypernova_batch(registry, request) - - response = handler(request) - + def hypernova_mapper(request, chunk): if not request.hypernova_batch.jobs: - return response + return chunk try: # Skip token replacement logic if explicitly flagged to if request.disable_hypernova_tween: - return response + return chunk except AttributeError: pass + transformed_chunk = chunk.decode('utf-8') with hypernova_token_replacement(request.hypernova_batch) as body: - body['content'] = response.text + body['content'] = transformed_chunk - response.text = body['content'] + transformed_chunk = body['content'] + return transformed_chunk.encode('utf-8') + def hypernova_tween(request): + request.hypernova_batch = configure_hypernova_batch(registry, request) + response = handler(request) + # Loop over all chunks in response.app_iter. Unlike accessing response.body + # or response.text directly, this avoids buffering and is important + # if our app_iter is a generator + # + # In cases where app_iter is a list (the default), this should work + # equivalently + response.app_iter = map(lambda chunk: hypernova_mapper(request, chunk), response.app_iter) return response return hypernova_tween diff --git a/tests/tweens_test.py b/tests/tweens_test.py index 0ac0287..7ba88dd 100644 --- a/tests/tweens_test.py +++ b/tests/tweens_test.py @@ -1,6 +1,7 @@ from unittest import mock import pytest +from pyramid.response import Response from pyramid_hypernova.rendering import RenderToken from pyramid_hypernova.tweens import hypernova_tween_factory @@ -14,9 +15,7 @@ def mock_setup(self): self.token = RenderToken('my-unique-id') mock_handler = mock.Mock() - mock_handler.return_value = mock.Mock( - text=str(self.token) - ) + mock_handler.return_value = Response(text=str(self.token)) self.mock_get_job_group_url = mock.Mock(return_value='http://localhost:8888/batch') @@ -47,6 +46,9 @@ def test_tween_replaces_tokens_when_disable_hypernova_tween_not_set(self): response = self.tween(self.mock_request) + # Access the response's body to ensure the batch request is made + response.body + self.mock_batch_request_factory.assert_called_once_with( get_job_group_url=self.mock_get_job_group_url, plugin_controller=mock.ANY, @@ -67,6 +69,10 @@ def test_tween_replaces_tokens_when_disable_hypernova_tween_set_false(self): json_encoder=self.mock_json_encoder, pyramid_request=self.mock_request, ) + + # Access the response's body to ensure the batch request is made + response.body + assert self.mock_batch_request_factory.return_value.submit.called assert response.text == '
REACT!
' @@ -75,6 +81,9 @@ def test_tween_replaces_tokens_when_disable_hypernova_tween_set_true(self): response = self.tween(self.mock_request) + # Access the response's body to ensure the batch request is made + response.body + self.mock_batch_request_factory.assert_called_once_with( get_job_group_url=self.mock_get_job_group_url, plugin_controller=mock.ANY, @@ -85,7 +94,7 @@ def test_tween_replaces_tokens_when_disable_hypernova_tween_set_true(self): assert response.text == str(self.token) def test_tween_returns_unmodified_response_if_no_jobs(self): - mock_response = mock.Mock( + mock_response = Response( body="I'm a binary file", ) mock_handler = mock.Mock() @@ -102,5 +111,8 @@ def test_tween_returns_unmodified_response_if_no_jobs(self): ): response = tween(self.mock_request) + # Access the response's body to ensure the batch request is made + response.body + assert not self.mock_batch_request_factory.return_value.submit.called assert response == mock_response