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_timeout 没有触发。

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

编辑

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

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() 以避免阻塞服务器.