How to test an aiohttp WebSocket in python

[mplatypus]

I am trying to figure out how to test a websocket in python. I cannot seem to get anything to work, in terms of actually getting the websocket to do anything.

The test:

@pytest.mark.asyncio
async def test_websocket(
    self,
    gateway_bot: gateway_bot_.GatewayBot,
    ongaku_client: Client,
    ongaku_session: Session,
    bot_user: OwnUser,
):
    ongaku_client._client_session = cs = mock.Mock()

    cs.ws_connect.return_value = ws = mock.MagicMock()
    ws.__aiter__.return_value = [
        aiohttp.WSMessage(aiohttp.WSMsgType.TEXT, orjson.dumps(payloads.READY_PAYLOAD).decode(), None), 
        aiohttp.WSMessage(aiohttp.WSMsgType.CLOSED, aiohttp.WSCloseCode.GOING_AWAY, "extra"),
        StopAsyncIteration
    ]

    with (
        mock.patch.object(gateway_bot, "get_me", return_value=bot_user),
        mock.patch.object(ongaku_session, "_handle_ws_message", return_value=None),
    ):
        await ongaku_session._websocket()

The method it is attempting to test:

async def _websocket(self) -> None:
    bot = self.app.get_me()

    _logger.log(
        TRACE_LEVEL,
        f"Attempting to start websocket connection to session {self.name}",
    )

    if not bot:
        if self._remaining_attempts > 0:
            self._status = session_.SessionStatus.NOT_CONNECTED

            _logger.warning(
                "Attempted fetching the bot, but failed as it does not exist."
            )
        else:
            self._status = session_.SessionStatus.FAILURE

        _logger.warning(
            "Attempted fetching the bot, but failed as it does not exist."
        )

        raise errors.SessionStartError

    self._websocket_headers = {
        "User-Id": str(int(bot.id)),
        "Client-Name": f"{bot.global_name if bot.global_name else 'unknown'}/{__version__}",
    }

    new_headers: typing.MutableMapping[str, typing.Any] = {}

    new_headers.update(self._websocket_headers)

    new_headers.update(self.auth_headers)

    while self._remaining_attempts >= 1:
        if self._remaining_attempts != self._attempts:
            await asyncio.sleep(2.5)
        self._remaining_attempts -= 1
        try:
            async with (
                self.client._get_client_session() as session,
                session.ws_connect(
                    self.base_uri + "/v4/websocket",
                    headers=new_headers,
                    autoclose=False,
                ) as ws,
            ):
                _logger.log(
                    TRACE_LEVEL,
                    f"Successfully made connection to session {self.name}",
                )
                self._status = session_.SessionStatus.CONNECTED
                async for msg in ws:
                    await self._handle_ws_message(msg)

        except Exception as e:
            _logger.warning(f"Websocket connection failure: {e}")
            self._status = session_.SessionStatus.NOT_CONNECTED

    else:
        _logger.critical("Server has no more attempts.")
        self._status = session_.SessionStatus.NOT_CONNECTED

Currently, it works fine, up until the following part:

async with (
    self.client._get_client_session() as session,
        session.ws_connect(
        self.base_uri + "/v4/websocket",
        headers=new_headers,
        autoclose=False,
    ) as ws,
):

where it just never actually returns (I have tested this by simply putting a raise Exception right after the ): on the end, which never actually returns.

What I expect to be able to do, is completely mock the response, by feeding it my own Websocket Messages, and then confirm that the _handle_ws_message method was called with the correct arguments (the two websocket messages)

[Dreamsorcerer]

Probably better to simulate the application, rather than mocking all the methods, something closer to what we have in our internal tests:
https://github.com/aio-libs/aiohttp/blob/master/tests/test_client_ws_functional.py