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

"""Stage-three PostgreSQL acceptance in dedicated databases only."""

import asyncio
import os
import subprocess
from pathlib import Path

from alembic import command
from alembic.config import Config
from cryptography.fernet import Fernet

NAME = "wq_research_stage3_test"
RESTORE = "wq_research_restore_stage3"
os.environ.update(
    DATABASE_URL=f"postgresql+asyncpg://postgres:research-test-only@127.0.0.1:18436/{NAME}",
    ADMIN_PASSWORD="research-acceptance-only",
    ENCRYPTION_KEY=Fernet.generate_key().decode(),
)


def docker(*args, **kwargs):
    return subprocess.run(["docker", "exec", "-i", "wq-research-acceptance-pg", *args], check=True, **kwargs)


async def acceptance():
    from unittest.mock import patch

    import httpx
    from sqlalchemy import select

    from app.config import Settings
    from app.main import create_app
    from app.models import ResearchInputSnapshot
    from app.research.workspace_contracts import TemplateSpec
    from tests.test_ai import configure
    from tests.test_backtests import setup
    from tests.test_research_flows import begin, get, test_fixed_two_rounds_and_idempotent_start
    from tests.test_research_workspace import template

    app = create_app(Settings(_env_file=None, enable_runner=False, public_origin="http://testserver"))
    async with app.router.lifespan_context(app):
        async with httpx.AsyncClient(
            transport=httpx.ASGITransport(app), base_url="http://testserver", headers={"X-WQ-Request": "1"}
        ) as client:
            assert (
                await client.post(
                    "/api/v1/auth/login", json={"username": "admin", "password": "research-acceptance-only"}
                )
            ).status_code == 200
            # Configure a deterministic model; no provider or real platform network.
            from tests.ai_fake import fake_model

            app.state.ai.model_factory = fake_model
            await configure(app, client)
            platform, lane = await setup(app)
            calls = []

            async def model(ai, context, output_type, revision):
                calls.append(context)
                value = template()
                value["expression"] = f"rank({{field}}) + {len(calls)}"
                return TemplateSpec.model_validate(value), {
                    "model": "fixture",
                    "revision": revision,
                    "usage": {"requests": 1},
                }

            async with app.state.sessions() as db:
                fixed = await db.scalar(select(ResearchInputSnapshot))
            body = {
                "request_id": "finite-run",
                "name": "PG 有限研究",
                "input_ids": [fixed.id],
                "hypothesis": "排名稳定性",
                "settings": {"region": "USA", "universe": "TOP3000", "delay": 1},
                "budget": {"max_rounds": 2, "max_simulations": 4, "max_model_calls": 3},
                "batch_candidates": 2,
            }
            with patch("app.research.runtime.request_model", model):
                await test_fixed_two_rounds_and_idempotent_start(app, client, (body, platform, lane, calls))
                short = {
                    **body,
                    "request_id": "concurrent-budget",
                    "budget": {"max_rounds": 1, "max_simulations": 1, "max_model_calls": 2},
                }
                first, second = await asyncio.gather(begin(client, short), begin(client, short))
                assert first["id"] == second["id"]
                from app.research.runtime import ResearchRuntime

                another = ResearchRuntime(app.state.sessions, app.state.ai, app.state.runner)
                for _ in range(7):
                    await asyncio.gather(
                        app.state.research.advance(first["id"]), another.advance(first["id"])
                    )
                result = await get(client, first["id"])
                assert result["status"] == "budget_exhausted" and result["simulations_used"] == 0
                assert result["model_calls_used"] == 1
    print("PASS PostgreSQL: two-round execution, idempotent starts, concurrent reservations and budget gate")


if __name__ == "__main__":
    docker("createdb", "-U", "postgres", NAME)
    with Path("/tmp/wq-research-stage2.dump").open("rb") as source:
        docker("pg_restore", "-U", "postgres", "-d", NAME, stdin=source)
    config = Config("alembic.ini")
    command.upgrade(config, "0008")
    command.check(config)
    asyncio.run(acceptance())
    dump = Path("/tmp/wq-research-stage3.dump")
    with dump.open("wb") as output:
        docker("pg_dump", "-U", "postgres", "-Fc", NAME, stdout=output)
    docker("createdb", "-U", "postgres", RESTORE)
    with dump.open("rb") as source:
        docker("pg_restore", "-U", "postgres", "-d", RESTORE, stdin=source)
    query = "SELECT (SELECT count(*) FROM research_revisions),(SELECT count(*) FROM research_flow_runs),(SELECT count(*) FROM research_step_runs),(SELECT note FROM research WHERE alpha_id='OLD_RESEARCH')"
    a = docker("psql", "-U", "postgres", "-d", NAME, "-Atc", query, capture_output=True).stdout
    b = docker("psql", "-U", "postgres", "-d", RESTORE, "-Atc", query, capture_output=True).stdout
    assert a == b
    print(
        "PASS PostgreSQL 17: 0007 → 0008 and pg_dump/pg_restore preserve flow budgets, steps and old research notes"
    )