120 lines
3.7 KiB
Python
120 lines
3.7 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
from collections.abc import AsyncIterator, Iterator
|
|
|
|
import pytest
|
|
import pytest_asyncio
|
|
from alembic.config import Config
|
|
from httpx import ASGITransport, AsyncClient
|
|
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
|
|
from testcontainers.postgres import PostgresContainer
|
|
|
|
from alembic import command
|
|
|
|
os.environ.setdefault("MONLET_AUTH_TOKEN", "test-token")
|
|
os.environ.setdefault("MONLET_ENABLE_DETECTOR", "false")
|
|
os.environ.setdefault("MONLET_ENABLE_NOTIFIER_WORKER", "false")
|
|
|
|
from monlet_server import db as db_module # noqa: E402
|
|
from monlet_server.main import create_app # noqa: E402
|
|
from monlet_server.settings import get_settings, reset_settings_cache # noqa: E402
|
|
|
|
HERE = os.path.dirname(__file__)
|
|
SERVER_ROOT = os.path.dirname(HERE)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def event_loop() -> Iterator[asyncio.AbstractEventLoop]:
|
|
loop = asyncio.new_event_loop()
|
|
yield loop
|
|
loop.close()
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def pg_container() -> Iterator[PostgresContainer]:
|
|
with PostgresContainer("postgres:16-alpine") as pg:
|
|
yield pg
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def database_url(pg_container: PostgresContainer) -> str:
|
|
raw = pg_container.get_connection_url()
|
|
if raw.startswith("postgresql+psycopg2://"):
|
|
async_url = raw.replace("postgresql+psycopg2://", "postgresql+asyncpg://")
|
|
elif raw.startswith("postgresql://"):
|
|
async_url = raw.replace("postgresql://", "postgresql+asyncpg://", 1)
|
|
else:
|
|
async_url = raw
|
|
return async_url
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def _apply_migrations(database_url: str) -> Iterator[None]:
|
|
os.environ["MONLET_DATABASE_URL"] = database_url
|
|
reset_settings_cache()
|
|
cfg = Config(os.path.join(SERVER_ROOT, "alembic.ini"))
|
|
cfg.set_main_option("script_location", os.path.join(SERVER_ROOT, "alembic"))
|
|
cfg.set_main_option("sqlalchemy.url", get_settings().sync_database_url)
|
|
command.upgrade(cfg, "head")
|
|
|
|
# Partitions for `events` are created at runtime, not by the migration.
|
|
async def _bootstrap() -> None:
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
|
|
from monlet_server.services.partitioning import ensure_event_partitions
|
|
|
|
eng = create_async_engine(database_url, future=True)
|
|
try:
|
|
await ensure_event_partitions(eng)
|
|
finally:
|
|
await eng.dispose()
|
|
|
|
asyncio.run(_bootstrap())
|
|
yield
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def engine(database_url: str):
|
|
eng = create_async_engine(database_url, future=True)
|
|
yield eng
|
|
await eng.dispose()
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def session(engine) -> AsyncIterator[AsyncSession]:
|
|
sm = async_sessionmaker(engine, expire_on_commit=False)
|
|
async with sm() as s:
|
|
yield s
|
|
|
|
|
|
@pytest_asyncio.fixture(autouse=True)
|
|
async def _truncate_tables(engine) -> AsyncIterator[None]:
|
|
yield
|
|
from sqlalchemy import text
|
|
|
|
async with engine.begin() as conn:
|
|
await conn.execute(
|
|
text(
|
|
"TRUNCATE TABLE notification_outbox, incidents, events, event_ingest_dedup, checks, agents RESTART IDENTITY CASCADE"
|
|
)
|
|
)
|
|
|
|
|
|
@pytest_asyncio.fixture
|
|
async def app_client(database_url: str) -> AsyncIterator[AsyncClient]:
|
|
reset_settings_cache()
|
|
await db_module.dispose_engine()
|
|
app = create_app()
|
|
async with app.router.lifespan_context(app):
|
|
transport = ASGITransport(app=app)
|
|
async with AsyncClient(transport=transport, base_url="http://test") as client:
|
|
yield client
|
|
await db_module.dispose_engine()
|
|
|
|
|
|
@pytest.fixture
|
|
def auth_headers() -> dict[str, str]:
|
|
return {"Authorization": "Bearer test-token"}
|