github twitter linkedin
Testing tornado websockets without third party clients

Recently, I built an encrypted chat service, which was based on tornado and ember.js. The project itself had grave security issues, so I shut it down, but while working on it I learned a few new things and testing websockets with tornado is one of them.

Most of the material out there for this suggests developing separate client based tests, which I didn’t want to do. Eventually, I figured out that tornado already provides all the utilities to do unit/integration tests for websockets.

First, we will need a websockets based echo server to test, lets call it ws.py. The websocket handler would be:

from tornado import web, websocket

class Echo(websocket.WebSocketHandler):

    # Open allows for any number arguments, unlike what pylint thinks.
    # pylint: disable=W0221
    def open(self):
        self.write_message('hello')

    def on_message(self, message):
        self.write_message(message)

    def on_close(self):
        self.write_message('bye')

Lets define an application which uses the above handler:

APP = web.Application([
    (r"/", Echo),
])

if __name__ == "__main__":
    APP.listen(5000)

Now, we will test the application out. Create a file, say test_ws.py:

from tornado import testing, httpserver, gen, websocket
from ws import APP

class TestChatHandler(testing.AsyncTestCase):
    pass

We use tornado’s testing wrapper for the integration it provides with the event loop. Lets tell unittest how to setup the tests:

class TestChatHandler(testing.AsyncTestCase):

    def setUp(self):
        super(TestChatHandler, self).setUp()
        server = httpserver.HTTPServer(APP)
        socket, self.port = testing.bind_unused_port()
        server.add_socket(socket)

We create a http server out of our application and get a socket bound to an unused port. We then ask the server to accept on the created socket. Don’t forget the super call, it ensures that the ioloop gets created. unittest will now ensure that a server and an ioloop is up and running before running tests.

Moving forward, we need to define a helper for creating a websocket connection to the server. Tornado websocket provides a handly websocket client. It can be created with websocket.websocket_connect.

    def _mk_connection(self):
        return websocket.websocket_connect(
            'ws://localhost:{}/'.format(self.port)
        )

We can write a simple test for this:

    @testing.gen_test
    def test_hello(self):
        c = yield self._mk_connection()
        # Get the initial hello from the server.
        response = yield c.read_message()
        # Make sure that we got a 'hello' not 'bye'
        self.assertEqual('hello', response)

testing.gen_test is a wrapper over tornado’s gen.coroutine. It runs the test synchronously under the ioloop that testing.AsyncTestCase creates in setUp. The test checks for the ‘hello’ message that we expect from the server on connection. yield makes sure that we for the response from the server. Note that if you write a yield c.read_message() when a message from server isn’t expected, the coroutine will keep waiting, eventually raising tornado.ioloop.TimeoutError (5 seconds by default). Great, we can write lot of tests using just what we have now.

The tests can be run via:

python -m tornado.testing discover

This could still be further improved. We need to yield and ignore the ‘hello’ message in every test, for every client. And in your application, it may be a more complicated handshake - possibly a few initial messages. Once you write a test for that handshake, it needn’t be re-written in every test. To avoid that, we will write an an abstraction over this:

    @gen.coroutine
    def _mk_client(self):
        c = yield self._mk_connection()

        # Discard the hello
        # This could be any initial handshake, which needs to be generalized
        # for most of the tests.
        _ = yield c.read_message()

        raise gen.Return(c)

_mk_client here is a method in which you could place all the boilerplate. The key point here is the exception gen.Return(c) we raise in the end. return with a value is allowed only after Python 3.3, so tornado.gen uses the value associated with this exception as the coroutine’s result.

With _mk_client available, we can write tests which only include the relevant code:

    @testing.gen_test
    def test_echo(self):
        # A client with the hello taken care of.
        c = yield self._mk_client()

        # Send a 'foo' to the server.
        c.write_message("foo")
        # Get the 'foo' back.
        response = yield c.read_message()
        # Make sure that we got a 'foo' back and not 'bar'.
        self.assertEqual('foo', response)

The application built out of this post is available as a gist.