-
Notifications
You must be signed in to change notification settings - Fork 31
/
websocket_server.py
60 lines (48 loc) · 1.92 KB
/
websocket_server.py
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
import threading
import pytest
from ws4py.server.wsgirefserver import WebSocketWSGIRequestHandler, WSGIServer
from ws4py.server.wsgiutils import WebSocketWSGIApplication
from wsgiref.simple_server import make_server
class WebsocketServer(threading.Thread):
""" Local websocket thread.
This class will allow us to test websocket requests locally, without
the need to mock any socket operations. It is a threading class, so
that we would be able to start and stop the server.
"""
def __init__(self, handler_class, host='127.0.0.1', port=0, **kwargs):
self._server = make_server(
host, port, server_class=WSGIServer,
handler_class=WebSocketWSGIRequestHandler,
app=WebSocketWSGIApplication(handler_cls=handler_class)
)
self._server.initialize_websockets_manager()
super(WebsocketServer, self).__init__(
name=self.__class__,
target=self._server.serve_forever
)
@property
def uri(self):
return 'ws://{}:{}'.format(*self._server.server_address)
def stop(self):
self._server.shutdown()
def wsserver_factory(server_handler_class):
""" websocket server factory. Creates a pytest factory which can be
later used in the tests.
Usage:
from websocket_server import wsserver_factory, WebSocket
# echo server implementation.
class MyServerHandler(WebSocket):
def handleMessage(self):
self.sendMessage(self.data)
wsserver = wsserver_factory(MyServerHandler)
def test_some_functionality(wsserver):
# you can now access the wsserver uri, like so:
print(wsserver.uri) # ws://localhost:some-port
"""
@pytest.fixture
def wsserver():
server = WebsocketServer(server_handler_class, 'localhost', 0)
server.start()
yield server
server.stop()
return wsserver