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

"""Stage-four 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_stage4_test"
RESTORE = "wq_research_restore_stage4"
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_quantflow import graph, launch, save
    from tests.test_research_flows import drive
    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,
            }
            body["request_id"] = "pg-quantflow"
            fixed_template = await save(client, "template", template())
            definition = graph("input", "expand", "backtest", "evaluate", "filter", "iterate")
            definition["nodes"][1]["config"] = {"asset_id": fixed_template["id"], "version": 1}
            definition["nodes"][4]["config"] = {"verdicts": ["pass", "review", "block"]}
            with patch("app.research.runtime.request_model", model):
                run = await launch(client, body, definition)
                result = await drive(app, client, run["id"], lane, ticks=60)
                assert result["status"] == "completed", result
                assert (
                    result["round"] == 2
                    and result["simulations_used"] == 4
                    and result["model_calls_used"] == 0
                )
                assert not calls and len(platform.posts) == 2
                await app.state.research.recover()
                await app.state.research.advance(run["id"])
                assert len(platform.posts) == 2
                print("PASS PostgreSQL: versioned native graph, filtering, bounded iteration and replay")


if __name__ == "__main__":
    docker("createdb", "-U", "postgres", NAME)
    with Path("/tmp/wq-research-stage3.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-stage4.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 md5(string_agg(row_to_json(t)::text, '' ORDER BY id)) FROM research_flow_runs t),(SELECT md5(string_agg(row_to_json(t)::text, '' ORDER BY id)) FROM research_step_runs t),(SELECT count(*) FROM research_revisions),(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: head 0008 unchanged, pg_dump/pg_restore preserves complete run/step snapshots and old research notes"
    )