234 lines
11 KiB
Python
234 lines
11 KiB
Python
"""Общие fixtures для интеграционных тестов, запускаемых против реального экземпляра Postgres
|
||
из `deploy/docker-compose.yml`.
|
||
|
||
Каждый тест запускается внутри внешней транзакции, которая никогда не коммитится; ORM
|
||
`AsyncSession` присоединяется к ней с помощью `join_transaction_mode="create_savepoint"` так,
|
||
что даже вызовы `session.commit()` (которые вызывают IntegrityError при нарушении
|
||
constraint) влияют только на savepoint и полностью отменяются внешним rollback в teardown.
|
||
"""
|
||
|
||
import asyncio
|
||
import json
|
||
from collections.abc import AsyncGenerator
|
||
from typing import Any, cast
|
||
|
||
import httpx
|
||
import pytest_asyncio
|
||
from fastapi import FastAPI
|
||
from httpx import ASGITransport
|
||
from sqlalchemy import delete, select
|
||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
|
||
from starlette.types import Message, Scope
|
||
|
||
from core.db import engine, get_session
|
||
from core.redis import redis_client
|
||
from main import create_app
|
||
from models.instance_setting import InstanceSetting
|
||
|
||
|
||
@pytest_asyncio.fixture(autouse=True)
|
||
async def _reset_rate_limits() -> AsyncGenerator[None, None]:
|
||
"""Сбросить счётчики rate limit (Redis реальный, общий на все тесты) до и после теста.
|
||
|
||
Без этого параллельные/последовательные тесты публичных эндпоинтов
|
||
(`resolve`, `guest-join`) делили бы один и тот же счётчик по IP тестового
|
||
клиента и мешали друг другу (см. `core/rate_limit.py`).
|
||
"""
|
||
await _delete_rate_limit_keys()
|
||
yield
|
||
await _delete_rate_limit_keys()
|
||
|
||
|
||
async def _delete_rate_limit_keys() -> None:
|
||
keys = [key async for key in redis_client.scan_iter(match="rate_limit:*")]
|
||
if keys:
|
||
await redis_client.delete(*keys)
|
||
|
||
|
||
@pytest_asyncio.fixture(autouse=True)
|
||
async def _preserve_instance_settings() -> AsyncGenerator[None, None]:
|
||
"""Гарантировать, что тест не оставляет следов в `instance_settings` общей dev-БД.
|
||
|
||
Строки `instance_settings` — живые настройки dev-инстанса (тоггл чата,
|
||
домен регистрации и т.п.), а не тестовые данные: закоммиченная тестом
|
||
правка молча меняет поведение dev-окружения (реально воспроизводилось —
|
||
после прогонов тестов чат в dev-инстансе оказался выключен).
|
||
Снимок закоммиченного состояния снимается отдельным подключением (мимо
|
||
savepoint-транзакции теста, см. докстринг модуля), после теста таблица
|
||
приводится к снимку: появившиеся ключи удаляются, изменённые и пропавшие —
|
||
восстанавливаются. Для корректных savepoint-тестов это no-op ценой одного
|
||
SELECT — страховка на случай любой записи мимо savepoint-сессии.
|
||
"""
|
||
before = await _load_committed_instance_settings()
|
||
yield
|
||
after = await _load_committed_instance_settings()
|
||
if after == before:
|
||
return
|
||
async with engine.connect() as connection:
|
||
extra_keys = after.keys() - before.keys()
|
||
if extra_keys:
|
||
await connection.execute(
|
||
delete(InstanceSetting).where(InstanceSetting.key.in_(extra_keys))
|
||
)
|
||
for key, value in before.items():
|
||
if after.get(key) != value:
|
||
await connection.execute(
|
||
pg_insert(InstanceSetting)
|
||
.values(key=key, value=value)
|
||
.on_conflict_do_update(index_elements=["key"], set_={"value": value})
|
||
)
|
||
await connection.commit()
|
||
|
||
|
||
async def _load_committed_instance_settings() -> dict[str, Any]:
|
||
"""Прочитать закоммиченные строки `instance_settings` отдельным подключением."""
|
||
async with engine.connect() as connection:
|
||
result = await connection.execute(select(InstanceSetting.key, InstanceSetting.value))
|
||
return {key: value for key, value in result.all()}
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def db_connection() -> AsyncGenerator[AsyncConnection, None]:
|
||
async with engine.connect() as connection:
|
||
trans = await connection.begin()
|
||
try:
|
||
yield connection
|
||
finally:
|
||
await trans.rollback()
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def db_session(db_connection: AsyncConnection) -> AsyncGenerator[AsyncSession, None]:
|
||
session = AsyncSession(
|
||
bind=db_connection,
|
||
join_transaction_mode="create_savepoint",
|
||
expire_on_commit=False,
|
||
)
|
||
try:
|
||
yield session
|
||
finally:
|
||
await session.close()
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def app(db_session: AsyncSession) -> AsyncGenerator[FastAPI, None]:
|
||
"""Экземпляр FastAPI-приложения с `get_session`, подменённым на тестовую (savepoint) сессию."""
|
||
application = create_app()
|
||
|
||
async def _override_get_session() -> AsyncGenerator[AsyncSession, None]:
|
||
yield db_session
|
||
|
||
application.dependency_overrides[get_session] = _override_get_session
|
||
yield application
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def client(app: FastAPI) -> AsyncGenerator[httpx.AsyncClient, None]:
|
||
"""Асинхронный HTTP-клиент поверх приложения.
|
||
|
||
База `https://test` (а не `http://`) нужна, чтобы httpx сохранял в своём
|
||
cookie jar httpOnly Secure cookie с refresh-токеном между запросами.
|
||
"""
|
||
transport = ASGITransport(app=app)
|
||
async with httpx.AsyncClient(transport=transport, base_url="https://test") as ac:
|
||
yield ac
|
||
|
||
|
||
# Таймаут ожидания ответа приложения в WS-тестах — на порядок больше
|
||
# `AUTH_TIMEOUT_SECONDS` эндпоинта не нужен, реальный ответ приходит мгновенно.
|
||
_RECEIVE_TIMEOUT_S = 5.0
|
||
|
||
|
||
class ASGIWebSocketSession:
|
||
"""Минимальный in-process ASGI websocket-клиент для тестов.
|
||
|
||
`httpx.AsyncClient`/`ASGITransport` не поддерживают websocket-соединения,
|
||
а `starlette.testclient.TestClient` гоняет ASGI-приложение в отдельном
|
||
потоке со своим event loop — тестовая `db_session` (см. выше) привязана к
|
||
`AsyncConnection` ТЕКУЩЕГО loop, и обращение к ней из другого loop роняет
|
||
asyncpg `RuntimeError: Future attached to a different loop`. Этот класс
|
||
драйвит `app(scope, receive, send)` напрямую в текущем event loop через
|
||
пару `asyncio.Queue`, эмулируя протокол ASGI websocket (см.
|
||
`starlette/websockets.py`).
|
||
"""
|
||
|
||
def __init__(self, app: FastAPI, path: str) -> None:
|
||
self._to_app: asyncio.Queue[Message] = asyncio.Queue()
|
||
self._from_app: asyncio.Queue[Message] = asyncio.Queue()
|
||
scope: Scope = {
|
||
"type": "websocket",
|
||
"asgi": {"version": "3.0", "spec_version": "2.3"},
|
||
"http_version": "1.1",
|
||
"scheme": "ws",
|
||
"path": path,
|
||
"raw_path": path.encode(),
|
||
"query_string": b"",
|
||
"headers": [],
|
||
"client": ("test-client", 12345),
|
||
"server": ("test-server", 80),
|
||
"subprotocols": [],
|
||
}
|
||
self._task = asyncio.create_task(app(scope, self._receive, self._send))
|
||
|
||
async def _receive(self) -> Message:
|
||
return await self._to_app.get()
|
||
|
||
async def _send(self, message: Message) -> None:
|
||
await self._from_app.put(message)
|
||
|
||
async def connect(self) -> Message:
|
||
"""Отправить `websocket.connect` и дождаться ответа (`accept`/`close`)."""
|
||
await self._to_app.put({"type": "websocket.connect"})
|
||
return await self._receive_from_app()
|
||
|
||
async def send_json(self, data: dict[str, Any]) -> None:
|
||
await self._to_app.put({"type": "websocket.receive", "text": json.dumps(data)})
|
||
|
||
async def receive_json(self) -> dict[str, Any]:
|
||
message = await self._receive_from_app()
|
||
assert message["type"] == "websocket.send", message
|
||
return cast(dict[str, Any], json.loads(message["text"]))
|
||
|
||
async def receive_close(self) -> int:
|
||
message = await self._receive_from_app()
|
||
assert message["type"] == "websocket.close", message
|
||
return int(message["code"])
|
||
|
||
async def _receive_from_app(self) -> Message:
|
||
# Фиксированный таймаут ожидания ответа приложения — тестовый
|
||
# хелпер не даёт вызывающей стороне переопределить его параметром
|
||
# (см. ASYNC109: явный параметр `timeout` у async-функции — плохая
|
||
# практика, вместо этого таймаут задаётся здесь одним местом).
|
||
async with asyncio.timeout(_RECEIVE_TIMEOUT_S):
|
||
return await self._from_app.get()
|
||
|
||
async def aclose(self) -> None:
|
||
"""Эмулировать разрыв соединения клиентом и дождаться завершения ASGI-приложения."""
|
||
if not self._task.done():
|
||
await self._to_app.put({"type": "websocket.disconnect", "code": 1000})
|
||
try:
|
||
await asyncio.wait_for(self._task, timeout=5.0)
|
||
except Exception: # noqa: BLE001 — best-effort teardown в тестах
|
||
self._task.cancel()
|
||
|
||
|
||
@pytest_asyncio.fixture
|
||
async def ws_client(app: FastAPI) -> AsyncGenerator[Any, None]:
|
||
"""Фабрика in-process WS-клиентов (`ASGIWebSocketSession`) поверх текущего `app`.
|
||
|
||
Возвращает callable `(path) -> ASGIWebSocketSession`; открытые сессии
|
||
закрываются автоматически по завершении теста.
|
||
"""
|
||
sessions: list[ASGIWebSocketSession] = []
|
||
|
||
def _factory(path: str) -> ASGIWebSocketSession:
|
||
session = ASGIWebSocketSession(app, path)
|
||
sessions.append(session)
|
||
return session
|
||
|
||
yield _factory
|
||
|
||
for session in sessions:
|
||
await session.aclose()
|