Flask-SocketIO单元测试中创建连接时添加HTTP Header

2019独角兽企业重金招聘Python工程师标准>>>

在使用flask-socketio中，有时候需要在创建连接时检查HTTP Header中是否有相关的信息。但在flask-socketio的单元测试工具中，并没有方法添加HTTP Header，但可以对框架进行简单的修改，就可以了。

改动的地方如下：

# 1. flask_socketio/__init__.py
def test_client(self, app, namespace=None, headers=None):
        """Return a simple SocketIO client that can be used for unit tests."""
        return SocketIOTestClient(app, self, namespace, headers)

# flask_socketio/test_client.py

    def __init__(self, app, socketio, namespace=None, headers=None):
        def _mock_send_packet(sid, pkt):
            if pkt.packet_type == packet.EVENT or \
                    pkt.packet_type == packet.BINARY_EVENT:
                if sid not in self.queue:
                    self.queue[sid] = []
                if pkt.data[0] == 'message' or pkt.data[0] == 'json':
                    self.queue[sid].append({'name': pkt.data[0],
                                            'args': pkt.data[1],
                                            'namespace': pkt.namespace or '/'})
                else:
                    self.queue[sid].append({'name': pkt.data[0],
                                            'args': pkt.data[1:],
                                            'namespace': pkt.namespace or '/'})
            elif pkt.packet_type == packet.ACK or \
                    pkt.packet_type == packet.BINARY_ACK:
                self.ack = {'args': pkt.data,
                            'namespace': pkt.namespace or '/'}

        self.app = app
        self.sid = uuid.uuid4().hex
        self.queue[self.sid] = []
        self.callback_counter = 0
        self.socketio = socketio
        socketio.server._send_packet = _mock_send_packet
        socketio.server.environ[self.sid] = {}
        if isinstance(socketio.server.manager, PubSubManager):
            raise RuntimeError('Test client cannot be used with a message '
                               'queue. Disable the queue on your test '
                               'configuration.')
        socketio.server.manager.initialize()
        self.connect(namespace, headers)

    def connect(self, namespace=None, headers=None):
        """Connect the client.

        :param namespace: The namespace for the client. If not provided, the
                          client connects to the server on the global
                          namespace.

        Note that it is usually not necessary to explicitly call this method,
        since a connection is automatically established when an instance of
        this class is created. An example where it this method would be useful
        is when the application accepts multiple namespace connections.
        """
        environ = EnvironBuilder('/socket.io').get_environ()
        environ['flask.app'] = self.app

        if isinstance(headers, dict):
            for k, v in headers.items():
                environ['HTTP_' + str(k).upper()] = str(v)

        self.socketio.server._handle_eio_connect(self.sid, environ)
        if namespace is not None and namespace != '/':
            pkt = packet.Packet(packet.CONNECT, namespace=namespace)
            with self.app.app_context():
                self.socketio.server._handle_eio_message(self.sid,
                                                         pkt.encode())

单元测试的代码：

import unittest
from flask_socketio import test_client
from test_app import app, socket_io


class FlaskTestCase(unittest.TestCase):

    def setUp(self):
        self.app = app
        self.app.config['TESTING'] = True

        self.socket_io_client = socket_io.test_client(self.app, headers={"self_header": "test_header_content"})

    def test_socket_io(self):
        pass

    def tearDown(self):
        pass

if __name__ == '__main__':
    unittest.main()

原理是根据uwsgi的PEP333定义，environ中以HTTP_开头的变量，是客户端提供的HTTP请求头。

转载于:https://my.oschina.net/jianming/blog/1600717