websockets
websockets copied to clipboard
Testing against a running websockets server
Sorry if this is a trivial question, but I've spend a lot of time on it by now.
I wrote a server that does wesocket connection and session management, and I am struggling to find a way to test it properly.
I'm starting the server like this:
async with websockets.serve(dispatcher.start, server_config['host'], server_config['port'],
process_request=health_check):
await asyncio.Future()
"Dispatcher" handles a few things like extracting locale and session ID from the HTTP Request, and then retrieves a session, and attaches it to a Manager object:
class Dispatcher:
def __init__(self, engine_config, session_manager, logger):
...
async def destroy(self):
if self.websocket.open:
self.websocket.close()
#
await self.incoming_message_task
async def start(self, websocket):
# cookie is problematic: https://websockets.readthedocs.io/en/stable/topics/authentication.html#sending-credentials
cookie = websocket.request_headers.get("sid", None)
locale = websocket.request_headers.get("locale", 'en')
# TODO: message validation
# start engine for this session
engine = DialogEngine(
self.logger,
self.engine_config['deployment_name'],
self.engine_config['endpoint'],
plugins=self.engine_config['plugins'])
# get or create session
session = self.session_manager.get_session(cookie)
# write back the session cookie
await websocket.send(
Message(AUTH, {"cookie": session.cookie.output(header="")}).serialize()
)
# create manager
manager = AllyManager(engine, session, self.logger, websocket)
await manager.start()
Now I want to test the Dispatcher by setting up a server in an async fixture and test it with a few edge cases like reconnect handling etc.
I've been through a couple iterations, currently on this:
@pytest_asyncio.fixture
@patch('allylib.dialog.engine.DialogEngine', 'turn', 'I did something')
async def server():
# set env vars to avoid errors, todo: use fixture if you ever find out how to
os.environ["AZURE_OPENAI_API_KEY"] = "xxx"
logger = logging.getLogger()
dialog_history_manager = DialogHistoryManager(logger, session_store_type='local')
engine_config = {}
dispatcher = Dispatcher(engine_config, dialog_history_manager, logger)
async with websockets.serve(dispatcher.start, '0.0.0.0', 1337):
yield
await asyncio.Future()
When running a test against it, I get [WinError 1225] The remote computer refused the network connection:
@pytest.mark.asyncio
async def test_run_server(server):
async with websockets.connect('ws://localhost:1337') as ws:
await ws.send('test')
Is there something I'm doing wrong (well, probably), is there a best practice to do testing like this?
Do I have to wrap the server into an asyncio task and dispatch it like that?
I found a really stupid way to accomplish this by abusing the yield mechanism in pytest-asyncio:
@pytest.fixture
async def server():
...
dispatcher = Dispatcher(engine_config, dialog_history_manager, logger)
async with websockets.serve(dispatcher.start, 'localhost', 1337):
try:
yield None
finally:
print('Destroying server')
and using it like this:
@pytest.mark.asyncio
async def test_run_server(server):
async for x in server:
async with connect('ws://localhost:1337') as ws:
await ws.send(json.dumps({'type': 'clientTurn', 'ask': 'What is my age again?'}))
while True:
r = await ws.recv()
print(r)
But I'm sure there must be a better way, right?
Hello, sorry for getting to your question so late — I was unable to find time for this project during the first half of the year.
Your question is half about pytest and half about websockets.
I was never a proficient pytest hacker so I won't be able to help much on that front.
About websockets, here's how I'm testing against a real websockets server in websockets' own test suite:
https://github.com/python-websockets/websockets/blob/e35c15a2a70c347f6f7a3e503ff1181ac35e1298/tests/asyncio/test_client.py#L16-L21
And here's the implementation of the run_server
context manager:
https://github.com/python-websockets/websockets/blob/e35c15a2a70c347f6f7a3e503ff1181ac35e1298/tests/asyncio/server.py#L41-L44
This looks very similar to what you're trying to implement. Your solution appears similar too, even if the async for x in server:
leaves me a bit puzzled.
Hope this helps!
On thing that could help: if you connect to port 0, the OS will allocate a random available port. This can help isolating tests from one another. Then you need to find which port was allocated so you can connect to it. Here's how I'm doing it:
https://github.com/python-websockets/websockets/blob/e35c15a2a70c347f6f7a3e503ff1181ac35e1298/tests/asyncio/client.py#L15-L27
https://github.com/python-websockets/websockets/blob/e35c15a2a70c347f6f7a3e503ff1181ac35e1298/tests/asyncio/server.py#L8-L12