Merge pull request #35 from Yelp/app-iter-loop

Loop over app_iter rather than accessing response.text

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

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 == '<div>REACT!</div>'
 
@@ -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