Fix a bug with IOStream.read_until_close with a streaming_callback.

bdarnell · bdarnell · commit 722731d571a7 · 2012-11-22T13:31:48.000-05:00
If the underlying socket was already closed the buffered data
would be passed to the final callback rather than the streaming callback.

Conflicts:
	website/sphinx/releases/next.rst
diff --git a/tornado/iostream.py b/tornado/iostream.py
@@ -176,8 +176,14 @@ def read_until_close(self, callback, streaming_callback=None):
         a ``streaming_callback`` is not used.
         """
         self._set_read_callback(callback)
+        self._streaming_callback = stack_context.wrap(streaming_callback)
         if self.closed():
-            self._run_callback(callback, self._consume(self._read_buffer_size))
+            if self._streaming_callback is not None:
+                self._run_callback(self._streaming_callback,
+                                   self._consume(self._read_buffer_size))
+            self._run_callback(self._read_callback,
+                               self._consume(self._read_buffer_size))
+            self._streaming_callback = None
             self._read_callback = None
             return
         self._read_until_close = True
diff --git a/tornado/test/iostream_test.py b/tornado/test/iostream_test.py
@@ -296,6 +296,42 @@ def test_close_buffered_data(self):
             server.close()
             client.close()
 
+    def test_read_until_close_after_close(self):
+        # Similar to test_delayed_close_callback, but read_until_close takes
+        # a separate code path so test it separately.
+        server, client = self.make_iostream_pair()
+        client.set_close_callback(self.stop)
+        try:
+            server.write(b("1234"))
+            server.close()
+            self.wait()
+            client.read_until_close(self.stop)
+            data = self.wait()
+            self.assertEqual(data, b("1234"))
+        finally:
+            server.close()
+            client.close()
+
+    def test_streaming_read_until_close_after_close(self):
+        # Same as the preceding test but with a streaming_callback.
+        # All data should go through the streaming callback,
+        # and the final read callback just gets an empty string.
+        server, client = self.make_iostream_pair()
+        client.set_close_callback(self.stop)
+        try:
+            server.write(b("1234"))
+            server.close()
+            self.wait()
+            streaming_data = []
+            client.read_until_close(self.stop,
+                                    streaming_callback=streaming_data.append)
+            data = self.wait()
+            self.assertEqual(b(''), data)
+            self.assertEqual(b('').join(streaming_data), b("1234"))
+        finally:
+            server.close()
+            client.close()
+
     def test_large_read_until(self):
         # Performance test: read_until used to have a quadratic component
         # so a read_until of 4MB would take 8 seconds; now it takes 0.25
