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

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

改动的地方如下:

1
2
3
4
# 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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
# 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())

单元测试的代码:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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请求头。