worldquant-alpha-system/backend/tests/backtest_postgres.py

"""Isolated PostgreSQL migration/concurrency acceptance. Never point at a personal database.

Run with DATABASE_URL ending in /wq_backtest_test, synthetic ADMIN_PASSWORD and
ENCRYPTION_KEY. Uses only mock WorldQuant HTTP and a disposable database.
"""

import asyncio
import os

import httpx
from alembic import command
from alembic.config import Config
from sqlalchemy import func, select

from app.alphas import upsert_alpha
from app.config import Settings
from app.db import create_database
from app.main import create_app
from app.models import BacktestEvent, BacktestResult, BacktestRun, Research, SimulationAttempt
from app.worldquant import WqClient
from tests.backtest_fake import Platform
from tests.test_backtests import candidate, preview, setup, start, tick


async def seed_old(settings):
    engine, sessions = create_database(settings.database_url)
    async with sessions.begin() as db:
        await upsert_alpha(
            db,
            {
                "id": "MIGRATION_ALPHA",
                "type": "REGULAR",
                "regular": {"code": "rank(close) + 0"},
                "settings": candidate()["settings"],
            },
        )
        await db.flush()
        research = await db.get(Research, "MIGRATION_ALPHA")
        research.note = "keep old research across upgrade and simulations"
    await engine.dispose()


async def acceptance(settings):
    fake = Platform()
    app = create_app(settings, WqClient(settings, transport=httpx.MockTransport(fake)))
    async with app.router.lifespan_context(app):
        async with httpx.AsyncClient(
            transport=httpx.ASGITransport(app=app),
            base_url="http://testserver",
            headers={"X-WQ-Request": "1"},
        ) as client:
            assert (
                await client.post(
                    "/api/v1/auth/login",
                    json={"username": "admin", "password": settings.admin_password.get_secret_value()},
                )
            ).status_code == 200
            fake, lane = await setup(app)
            fake.existing_alpha_ids = ["MIGRATION_ALPHA"]
            p = await preview(client, [candidate(0), candidate(0) | {"client_item_id": "repeat"}])
            a, b = await asyncio.gather(
                start(client, p, "concurrent-confirm"), start(client, p, "concurrent-confirm")
            )
            assert a["backtest_run_id"] == b["backtest_run_id"]
            rid = a["backtest_run_id"]
            for _ in range(5):
                await tick(lane)
            result = (await client.get(f"/api/v1/backtests/runs/{rid}")).json()
            assert result["status"] == "completed", result
            assert len(fake.posts) == 2
            async with app.state.sessions() as db:
                assert await db.scalar(select(func.count()).select_from(BacktestRun)) == 1
                assert await db.scalar(select(func.count()).select_from(BacktestResult)) == 2
                note = (await db.get(Research, "MIGRATION_ALPHA")).note
                assert note == "keep old research across upgrade and simulations"
                events = list(
                    await db.scalars(
                        select(BacktestEvent.seq)
                        .where(BacktestEvent.run_id == rid)
                        .order_by(BacktestEvent.seq)
                    )
                )
                assert events == list(range(1, len(events) + 1))
            # Leave an accepted run for a new application instance to recover.
            next_run = await start(client, await preview(client, [candidate(2)]), "restart")
            async with app.state.sessions() as db:
                aid = await db.scalar(
                    select(SimulationAttempt.id).where(
                        SimulationAttempt.run_id == next_run["backtest_run_id"]
                    )
                )
            await lane.step(aid)
            await lane.interrupt()
    replacement = create_app(settings, WqClient(settings, transport=httpx.MockTransport(fake)))
    async with replacement.router.lifespan_context(replacement):
        lane = replacement.state.runner.backtests
        await lane.start()
        await lane.stop()
        await lane.step(aid)
        async with replacement.state.sessions() as db:
            assert (await db.get(BacktestRun, next_run["backtest_run_id"])).status == "completed"
        assert len(fake.posts) == 3
    print(
        "PASS PostgreSQL: concurrent confirmation creates one run; two attempts share one Alpha safely; contiguous transactional events; research preserved; replacement application resumes accepted simulation without POST"
    )


def main():
    if not os.environ.get("DATABASE_URL", "").endswith("/wq_backtest_test"):
        raise SystemExit("Only an isolated wq_backtest_test database is allowed")
    settings = Settings(_env_file=None, enable_runner=False, public_origin="http://testserver")
    config = Config("alembic.ini")
    command.upgrade(config, "0002")
    asyncio.run(seed_old(settings))
    command.upgrade(config, "head")
    command.check(config)
    asyncio.run(acceptance(settings))


if __name__ == "__main__":
    main()