python-websockets / websockets

Library for building WebSocket servers and clients in Python
https://websockets.readthedocs.io/
BSD 3-Clause "New" or "Revised" License
5.16k stars 513 forks source link

Testing against a running websockets server #1442

Closed maxupp closed 1 month ago

maxupp commented 7 months ago

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?

maxupp commented 7 months ago

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?

aaugustin commented 1 month ago

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!

aaugustin commented 1 month ago

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