web: Use write() callable for output.

[jelmer/dulwich-libgit2.git] / dulwich / tests / test_web.py

diff --git a/dulwich/tests/test_web.py b/dulwich/tests/test_web.py
index 0ac31fbc33a6028194790c6d27a050834cd1a015..9e5010a79251626a5b77f9a31379fcbb356c5776 100644 (file)
--- a/dulwich/tests/test_web.py
+++ b/dulwich/tests/test_web.py
@@ -50,10 +50,12 @@ class WebTestCase(TestCase):
                                    handlers=self._handlers())
         self._status = None
         self._headers = []
+        self._output = StringIO()
 
     def _start_response(self, status, headers):
         self._status = status
         self._headers = list(headers)
+        return self._output.write
 
     def _handlers(self):
         return None
@@ -192,24 +194,30 @@ class SmartHandlersTestCase(WebTestCase):
         list(handle_service_request(self._req, 'backend', mat))
         self.assertEquals(HTTP_FORBIDDEN, self._status)
 
-    def test_handle_service_request(self):
+    def _run_handle_service_request(self, content_length=None):
         self._environ['wsgi.input'] = StringIO('foo')
+        if content_length is not None:
+            self._environ['CONTENT_LENGTH'] = content_length
         mat = re.search('.*', '/git-upload-pack')
-        output = ''.join(handle_service_request(self._req, 'backend', mat))
-        self.assertEqual('handled input: foo', output)
+        handler_output = ''.join(
+          handle_service_request(self._req, 'backend', mat))
+        write_output = self._output.getvalue()
+        # Ensure all output was written via the write callback.
+        self.assertEqual('', handler_output)
+        self.assertEqual('handled input: foo', write_output)
         response_type = 'application/x-git-upload-pack-response'
         self.assertTrue(('Content-Type', response_type) in self._headers)
         self.assertFalse(self._handler.advertise_refs)
         self.assertTrue(self._handler.stateless_rpc)
 
+    def test_handle_service_request(self):
+        self._run_handle_service_request()
+
     def test_handle_service_request_with_length(self):
-        self._environ['wsgi.input'] = StringIO('foobar')
-        self._environ['CONTENT_LENGTH'] = 3
-        mat = re.search('.*', '/git-upload-pack')
-        output = ''.join(handle_service_request(self._req, 'backend', mat))
-        self.assertEqual('handled input: foo', output)
-        response_type = 'application/x-git-upload-pack-response'
-        self.assertTrue(('Content-Type', response_type) in self._headers)
+        self._run_handle_service_request(content_length='3')
+
+    def test_handle_service_request_empty_length(self):
+        self._run_handle_service_request(content_length='')
 
     def test_get_info_refs_unknown(self):
         self._environ['QUERY_STRING'] = 'service=git-evil-handler'
@@ -221,11 +229,14 @@ class SmartHandlersTestCase(WebTestCase):
         self._environ['QUERY_STRING'] = 'service=git-upload-pack'
 
         mat = re.search('.*', '/git-upload-pack')
-        output = ''.join(get_info_refs(self._req, 'backend', mat))
+        handler_output = ''.join(get_info_refs(self._req, 'backend', mat))
+        write_output = self._output.getvalue()
         self.assertEquals(('001e# service=git-upload-pack\n'
                            '0000'
                            # input is ignored by the handler
-                           'handled input: '), output)
+                           'handled input: '), write_output)
+        # Ensure all output was written via the write callback.
+        self.assertEquals('', handler_output)
         self.assertTrue(self._handler.advertise_refs)
         self.assertTrue(self._handler.stateless_rpc)