Реализация и тестирование таймаута подключения WebSocket server

Я внедряю сервер WebSockets в Tornado 3.2. Клиент, подключающийся к серверу, не будет браузером.

Для случаев, когда между сервером и клиентом существует обратная связь, я хотел бы добавить max. время, в течение которого сервер будет ждать ответа клиента перед закрытием соединения.

Это примерно то, что я пытался сделать:

import datetime
import tornado

class WSHandler(WebSocketHandler):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.timeout = None

    def _close_on_timeout(self):
        if self.ws_connection:
            self.close()

    def open(self):
        initialize()

    def on_message(self, message):
        # Remove previous timeout, if one exists.
        if self.timeout:
            tornado.ioloop.IOLoop.instance().remove_timeout(self.timeout)
            self.timeout = None

        if is_last_message:
            self.write_message(message)
            self.close()
        else:
            # Add a new timeout.
            self.timeout = tornado.ioloop.IOLoop.instance().add_timeout(
                datetime.timedelta(milliseconds=1000), self._close_on_timeout)
            self.write_message(message)

Я что, недотепа и есть ли более простой способ сделать это? Я даже не могу, кажется, запланировать простой распечатайте инструкцию через add_timeout выше.

Мне также нужна помощь в тестировании этого. Вот что у меня есть до сих пор:

from tornado.websocket import websocket_connect
from tornado.testing import AsyncHTTPTestCase, gen_test
import time

class WSTests(AsyncHTTPTestCase):

    @gen_test
    def test_long_response(self):
        ws = yield websocket_connect('ws://address', io_loop=self.io_loop)

        # First round trip.
        ws.write_message('First message.')
        result = yield ws.read_message()
        self.assertEqual(result, 'First response.')

        # Wait longer than the timeout.
        # The test is in its own IOLoop, so a blocking sleep should be okay?
        time.sleep(1.1)

        # Expect either write or read to fail because of a closed socket.
        ws.write_message('Second message.')
        result = yield ws.read_message()

        self.assertNotEqual(result, 'Second response.')

У клиента нет проблем с записью и чтением из сокета. Это, вероятно, потому, что add_timeout не срабатывает.

Должен ли тест каким-то образом уступить, чтобы разрешить обратный вызов таймаута на сервере? Я бы подумал, что нет, поскольку врачи говорят, что тесты выполняются в их собственном IOLoop.

Edit

Это и есть работа версия, согласно предложениям Бена.

import datetime
import tornado

class WSHandler(WebSocketHandler):

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.timeout = None

    def _close_on_timeout(self):
        if self.ws_connection:
            self.close()

    def open(self):
        initialize()

    def on_message(self, message):
        # Remove previous timeout, if one exists.
        if self.timeout:
            tornado.ioloop.IOLoop.current().remove_timeout(self.timeout)
            self.timeout = None

        if is_last_message:
            self.write_message(message)
            self.close()
        else:
            # Add a new timeout.
            self.timeout = tornado.ioloop.IOLoop.current().add_timeout(
                datetime.timedelta(milliseconds=1000), self._close_on_timeout)
            self.write_message(message)

Испытание:

from tornado.websocket import websocket_connect
from tornado.testing import AsyncHTTPTestCase, gen_test
import time

class WSTests(AsyncHTTPTestCase):

    @gen_test
    def test_long_response(self):
        ws = yield websocket_connect('ws://address', io_loop=self.io_loop)

        # First round trip.
        ws.write_message('First message.')
        result = yield ws.read_message()
        self.assertEqual(result, 'First response.')

        # Wait a little more than the timeout.
        yield gen.Task(self.io_loop.add_timeout, datetime.timedelta(seconds=1.1))

        # Expect either write or read to fail because of a closed socket.
        ws.write_message('Second message.')
        result = yield ws.read_message()
        self.assertEqual(result, None)