Python 实现和测试WebSocket服务器连接超时

我正在Tornado 3.2中实现WebSockets服务器。连接到服务器的客户端将不是浏览器

对于服务器和客户端之间存在来回通信的情况，我想添加一个最长时间，服务器在关闭连接之前等待客户端响应

这大概就是我一直在尝试的：

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_超时没有触发

测试是否需要以某种方式产生，以允许服务器上的超时回调运行？我本以为不会，因为医生说测试是在他们自己的IOLoop中运行的

编辑

根据本的建议，这是工作版本

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)

---

在我看来，第一个示例中的超时处理代码是正确的

---

对于测试，每个测试用例都有自己的IOLoop，但是对于测试和它运行的任何其他操作，只有一个IOLoop，因此您必须在此处使用add_timeout而不是time.sleep（），以避免阻塞服务器。

Ey Ben，我知道这个问题很久以前就解决了，但我想与阅读本文的任何用户分享我为此制定的解决方案。 它基本上是基于您的，但它解决了外部服务的问题，外部服务可以使用组合而不是继承轻松集成到任何websocket中：

class TimeoutWebSocketService():
    _default_timeout_delta_ms = 10 * 60 * 1000  # 10 min

    def __init__(self, websocket, ioloop=None, timeout=None):
        # Timeout
        self.ioloop = ioloop or tornado.ioloop.IOLoop.current()
        self.websocket = websocket
        self._timeout = None
        self._timeout_delta_ms = timeout or TimeoutWebSocketService._default_timeout_delta_ms

    def _close_on_timeout(self):
        self._timeout = None
        if self.websocket.ws_connection:
            self.websocket.close()

    def refresh_timeout(self, timeout=None):
        timeout = timeout or self._timeout_delta_ms
        if timeout > 0:
            # Clean last timeout, if one exists
            self.clean_timeout()

            # Add a new timeout (must be None from clean).
            self._timeout = self.ioloop.add_timeout(
                datetime.timedelta(milliseconds=timeout), self._close_on_timeout)

    def clean_timeout(self):
        if self._timeout is not None:
            # Remove previous timeout, if one exists.
            self.ioloop.remove_timeout(self._timeout)
            self._timeout = None

为了使用该服务，创建一个新的TimeoutWebService实例（可选超时单位为ms，以及应该执行该服务的ioloop）并调用“刷新超时”方法来设置第一次超时或重置一个已经存在的超时即可，或“清除超时”以停止超时服务

class BaseWebSocketHandler(WebSocketHandler):
    def prepare(self):
        self.timeout_service = TimeoutWebSocketService(timeout=(1000*60))

        ## Optionally starts the service here 
        self.timeout_service.refresh_timeout()

        ## rest of prepare method 

    def on_message(self):
        self.timeout_service.refresh_timeout()

    def on_close(self):
        self.timeout_service.clean_timeout()

多亏了这种方法，您可以准确地控制何时以及在何种条件下重新启动超时，这可能因应用程序而异。例如，您可能只希望在用户满足X条件或消息是预期消息时刷新超时

---

我希望ppl喜欢这个解决方案

谢谢你，本。在上面添加了一个编辑以显示我正在尝试的内容。测试超时回调将运行，但处理程序中的超时回调仍不会触发。我是否正确启动IOLoop？如果您使用的是

@gen\u test

，它将为您启动/停止IOLoop，因此您应该删除更新测试的最后两行。在协同程序中，您可能希望使用

yield gen.Task（self.io\u loop.add\u timeout，timedelta（seconds=1.1））

而不是通过回调来添加超时。混合协同路由和回调可能会很棘手（在这里，你需要用

@gen.coroutine

修饰

\u check

，并使用

io\u循环。添加未来的计划）Ben，我很抱歉太密集了。我做了更改，果然，测试现在等待了正确的时间。但是，服务器的超时回调仍然没有被调用，测试可以在等待1.1秒后读取和写入套接字。有什么想法吗？你需要使用IOLoop.current（），而不是IOLoop.instance（）。谢谢你的工作版本。这帮了大忙。