Testing against a running websockets server
maxupp opened this issue · comments
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?