Application to demonstrate testing tornado websockets

test_ws.py

'''
Application to demonstrate testing tornado websockets.

Try it wiith: python -m tornado.testing discover
'''


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

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)

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

    @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)

    @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
    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)

ws.py

'''
Application to demonstrate testing tornado websockets
'''

from tornado import web, websocket

# We don't override `data_received`
# pylint: disable=W0223
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')


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

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