Add close method to websocket client; fix socket leaks in websocket_test.

These leaks are detected only in python 3.4 thanks to its improved
generator GC.

diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
index 7dc06c5d261ceb9cf739cf8a44ef2cb79962ea14..da9d780b1140f1a686cb93cd4d51167c5cb3687a 100644 (file)
--- a/tornado/test/websocket_test.py
+++ b/tornado/test/websocket_test.py
@@ -5,19 +5,24 @@ from tornado.testing import AsyncHTTPTestCase, gen_test, bind_unused_port, Expec
 from tornado.web import Application, RequestHandler
 from tornado.websocket import WebSocketHandler, websocket_connect, WebSocketError
 
+class TestWebSocketHandler(WebSocketHandler):
+    """Base class for testing handlers that exposes the on_close event.
 
-class EchoHandler(WebSocketHandler):
+    This allows for deterministic cleanup of the associated socket.
+    """
     def initialize(self, close_future):
         self.close_future = close_future
 
-    def on_message(self, message):
-        self.write_message(message, isinstance(message, bytes))
-
     def on_close(self):
         self.close_future.set_result(None)
 
 
-class HeaderHandler(WebSocketHandler):
+class EchoHandler(TestWebSocketHandler):
+    def on_message(self, message):
+        self.write_message(message, isinstance(message, bytes))
+
+
+class HeaderHandler(TestWebSocketHandler):
     def open(self):
         self.write_message(self.request.headers.get('X-Test', ''))
 
@@ -33,7 +38,7 @@ class WebSocketTest(AsyncHTTPTestCase):
         return Application([
             ('/echo', EchoHandler, dict(close_future=self.close_future)),
             ('/non_ws', NonWebSocketHandler),
-            ('/header', HeaderHandler),
+            ('/header', HeaderHandler, dict(close_future=self.close_future)),
         ])
 
     @gen_test
@@ -44,6 +49,8 @@ class WebSocketTest(AsyncHTTPTestCase):
         ws.write_message('hello')
         response = yield ws.read_message()
         self.assertEqual(response, 'hello')
+        ws.close()
+        yield self.close_future
 
     def test_websocket_callbacks(self):
         websocket_connect(
@@ -54,6 +61,8 @@ class WebSocketTest(AsyncHTTPTestCase):
         ws.read_message(self.stop)
         response = self.wait().result()
         self.assertEqual(response, 'hello')
+        ws.close()
+        yield self.close_future
 
     @gen_test
     def test_websocket_http_fail(self):
@@ -99,3 +108,5 @@ class WebSocketTest(AsyncHTTPTestCase):
                         headers={'X-Test': 'hello'}))
         response = yield ws.read_message()
         self.assertEqual(response, 'hello')
+        ws.close()
+        yield self.close_future

diff --git a/tornado/websocket.py b/tornado/websocket.py
index cd1ca7f51d59c9d7d618f9e78d071ef0dd7c8ccd..75b1d8f23072b4a7ed0a704150452958c493935b 100644 (file)
--- a/tornado/websocket.py
+++ b/tornado/websocket.py
@@ -793,6 +793,12 @@ class WebSocketClientConnection(simple_httpclient._HTTPConnection):
             io_loop, None, request, lambda: None, self._on_http_response,
             104857600, self.resolver)
 
+    def close(self):
+        """Closes the websocket connection."""
+        if self.protocol is not None:
+            self.protocol.close()
+            self.protocol = None
+
     def _on_close(self):
         self.on_message(None)
         self.resolver.close()