import asyncio

from fastapi.testclient import TestClient
from starlette.testclient import TestClient

from asynctestws import async_websocket_connect


async def test_async_ws(
    client: TestClient,
):
    with async_websocket_connect(
        client,
        "/api/ws/connect",
        receive_timeout=3,
    ) as websocket:
        await websocket.asend_json(
            {
                "type": "answer",
                "secret": 42,
            }
        )
        response = await websocket.areceive_json()
        assert response["reaction"] == "prosperity"


# @pytest.mark.timeout(10)  # May be a good thing to use
async def test_async_ws_loop_nones(
    client: TestClient,
):
    with async_websocket_connect(
        client,
        "/api/ws/connect",
        receive_timeout=3,
        raise_timeout=False,  # Return None in case of timeout
    ) as websocket:
        await websocket.asend_json(
            {
                "type": "subscribe",
                "topic": "1337",
            }
        )
        done = 0
        failed = 0
        while done < 10:
            response = await websocket.areceive_json()
            if response is None:
                failed += 1
                continue
            assert "value" in response
            done += 1
        assert done > failed


# @pytest.mark.timeout(10)  # May be a good thing to use
async def test_async_ws_loop_raise(
    client: TestClient,
):
    with async_websocket_connect(
        client,
        "/api/ws/connect",
        receive_timeout=3,
        raise_timeout=True,
    ) as websocket:
        await websocket.asend_json(
            {
                "type": "subscribe",
                "topic": "1337",
            }
        )
        done = 0
        failed = 0
        while done < 10:
            try:
                response = await websocket.areceive_json()
                assert "value" in response
                done += 1
            except asyncio.exceptions.TimeoutError:
                failed += 1
                continue
        assert done > failed