Make the websocket_connect error handling more comprehensive.

bdarnell · bdarnell · commit 1e69faa11594 · 2013-04-08T23:11:08.000-04:00
Now covers successful http responses that return a non-websocket body
and network errors that prevent a body from being returned.

Added a connect_timeout parameter to websocket_connect.
diff --git a/docs/releases/v3.0.1.rst b/docs/releases/v3.0.1.rst
@@ -12,6 +12,8 @@ Apr 8, 2013
   as a (broken) test by ``nose``.
 * Work around a bug in Ubuntu 13.04 betas involving an incomplete backport
   of the `ssl.match_hostname` function.
+* `tornado.websocket.websocket_connect` now fails cleanly when it attempts
+  to connect to a non-websocket url.
 * `tornado.testing.LogTrapTestCase` once again works with byte strings
   on Python 2.
 * The ``request`` attribute of `tornado.httpclient.HTTPResponse` is
diff --git a/tornado/test/websocket_test.py b/tornado/test/websocket_test.py
@@ -1,17 +1,24 @@
-from tornado.testing import AsyncHTTPTestCase, gen_test
-from tornado.web import Application
-from tornado.websocket import WebSocketHandler, websocket_connect
+from tornado.httpclient import HTTPError
+from tornado.log import gen_log
+from tornado.testing import AsyncHTTPTestCase, gen_test, bind_unused_port, ExpectLog
+from tornado.web import Application, RequestHandler
+from tornado.websocket import WebSocketHandler, websocket_connect, WebSocketError
 
 
 class EchoHandler(WebSocketHandler):
     def on_message(self, message):
         self.write_message(message, isinstance(message, bytes))
 
+class NonWebSocketHandler(RequestHandler):
+    def get(self):
+        self.write('ok')
+
 
 class WebSocketTest(AsyncHTTPTestCase):
     def get_app(self):
         return Application([
             ('/echo', EchoHandler),
+            ('/non_ws', NonWebSocketHandler),
         ])
 
     @gen_test
@@ -32,14 +39,30 @@ def test_websocket_callbacks(self):
         ws.read_message(self.stop)
         response = self.wait().result()
         self.assertEqual(response, 'hello')
-    
+
     @gen_test
-    def test_websocket_fail(self):
-        try:
-            ws = yield websocket_connect(
-                'ws://localhost:%d/no_websock' % self.get_http_port(),
+    def test_websocket_http_fail(self):
+        with self.assertRaises(HTTPError) as cm:
+            yield websocket_connect(
+                'ws://localhost:%d/notfound' % self.get_http_port(),
                 io_loop=self.io_loop)
-        except:
-            pass
-        else:
-            self.fail('Should\'ve caught an Exception')
+        self.assertEqual(cm.exception.code, 404)
+
+    @gen_test
+    def test_websocket_http_success(self):
+        with self.assertRaises(WebSocketError):
+            yield websocket_connect(
+                'ws://localhost:%d/non_ws' % self.get_http_port(),
+                io_loop=self.io_loop)
+
+    @gen_test
+    def test_websocket_network_fail(self):
+        sock, port = bind_unused_port()
+        sock.close()
+        with self.assertRaises(HTTPError) as cm:
+            with ExpectLog(gen_log, ".*"):
+                yield websocket_connect(
+                    'ws://localhost:%d/' % port,
+                    io_loop=self.io_loop,
+                    connect_timeout=0.01)
+        self.assertEqual(cm.exception.code, 599)
diff --git a/tornado/websocket.py b/tornado/websocket.py
@@ -46,6 +46,10 @@
     xrange = range  # py3
 
 
+class WebSocketError(Exception):
+    pass
+
+
 class WebSocketHandler(tornado.web.RequestHandler):
     """Subclass this class to create a basic WebSocket handler.
 
@@ -740,14 +744,19 @@ def __init__(self, io_loop, request):
         })
 
         super(WebSocketClientConnection, self).__init__(
-            io_loop, None, request, lambda: None, lambda response: None,
+            io_loop, None, request, lambda: None, self._on_http_response,
             104857600, Resolver(io_loop=io_loop))
 
     def _on_close(self):
         self.on_message(None)
 
-    def _on_body(self, body):
-        self.connect_future.set_exception(Exception('Could not connect.'))
+    def _on_http_response(self, response):
+        if not self.connect_future.done():
+            if response.error:
+                self.connect_future.set_exception(response.error)
+            else:
+                self.connect_future.set_exception(WebSocketError(
+                        "Non-websocket response"))
 
     def _handle_1xx(self, code):
         assert code == 101
@@ -798,15 +807,15 @@ def on_pong(self, data):
         pass
 
 
-def websocket_connect(url, io_loop=None, callback=None):
+def websocket_connect(url, io_loop=None, callback=None, connect_timeout=None):
     """Client-side websocket support.
 
     Takes a url and returns a Future whose result is a
     `WebSocketClientConnection`.
     """
     if io_loop is None:
         io_loop = IOLoop.current()
-    request = httpclient.HTTPRequest(url)
+    request = httpclient.HTTPRequest(url, connect_timeout=connect_timeout)
     request = httpclient._RequestProxy(
         request, httpclient.HTTPRequest._DEFAULTS)
     conn = WebSocketClientConnection(io_loop, request)
