python-websockets / websockets

Library for building WebSocket servers and clients in Python

Home Page:https://websockets.readthedocs.io/

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

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?