nonebot2/tests/test_driver.py
2022-08-14 19:41:00 +08:00

105 lines
2.9 KiB
Python

from typing import cast
import pytest
from nonebug import App
@pytest.mark.asyncio
@pytest.mark.parametrize(
"nonebug_init",
[
pytest.param({"driver": "nonebot.drivers.fastapi:Driver"}, id="fastapi"),
pytest.param({"driver": "nonebot.drivers.quart:Driver"}, id="quart"),
],
indirect=True,
)
async def test_reverse_driver(app: App):
import nonebot
from nonebot.exception import WebSocketClosed
from nonebot.drivers import (
URL,
Request,
Response,
WebSocket,
ReverseDriver,
HTTPServerSetup,
WebSocketServerSetup,
)
driver = cast(ReverseDriver, nonebot.get_driver())
async def _handle_http(request: Request) -> Response:
assert request.content in (b"test", "test")
return Response(200, content="test")
async def _handle_ws(ws: WebSocket) -> None:
await ws.accept()
data = await ws.receive()
assert data == "ping"
await ws.send("pong")
data = await ws.receive()
assert data == b"ping"
await ws.send(b"pong")
data = await ws.receive_text()
assert data == "ping"
await ws.send("pong")
data = await ws.receive_bytes()
assert data == b"ping"
await ws.send(b"pong")
with pytest.raises(WebSocketClosed):
await ws.receive()
http_setup = HTTPServerSetup(URL("/http_test"), "POST", "http_test", _handle_http)
driver.setup_http_server(http_setup)
ws_setup = WebSocketServerSetup(URL("/ws_test"), "ws_test", _handle_ws)
driver.setup_websocket_server(ws_setup)
async with app.test_server() as ctx:
client = ctx.get_client()
response = await client.post("/http_test", data="test")
assert response.status_code == 200
assert response.text == "test"
async with client.websocket_connect("/ws_test") as ws:
await ws.send_text("ping")
assert await ws.receive_text() == "pong"
await ws.send_bytes(b"ping")
assert await ws.receive_bytes() == b"pong"
await ws.send_text("ping")
assert await ws.receive_text() == "pong"
await ws.send_bytes(b"ping")
assert await ws.receive_bytes() == b"pong"
await ws.close()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"nonebug_init, driver_type",
[
pytest.param(
{"driver": "nonebot.drivers.fastapi:Driver+nonebot.drivers.aiohttp:Mixin"},
"fastapi+aiohttp",
id="fastapi+aiohttp",
),
pytest.param(
{"driver": "~httpx:Driver+~websockets"},
"block_driver+httpx+websockets",
id="httpx+websockets",
),
],
indirect=["nonebug_init"],
)
async def test_combine_driver(app: App, driver_type: str):
import nonebot
driver = nonebot.get_driver()
assert driver.type == driver_type