在使用flask-socketio中，有時(shí)候需要在創(chuàng)建連接時(shí)檢查HTTP Header中是否有相關(guān)的信息种吸。但在flask-socketio的單元測(cè)試工具中，并沒(méi)有方法添加HTTP Header，但可以對(duì)框架進(jìn)行簡(jiǎn)單的修改吃型，就可以了。

改動(dòng)的地方如下：

# 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())

單元測(cè)試的代碼：

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

原理是根據(jù)uwsgi的PEP333定義僚楞，environ中以HTTP_開(kāi)頭的變量勤晚，是客戶端提供的HTTP請(qǐng)求頭。