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

"""Stage-two 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_stage2_test"
RESTORE = "wq_research_restore_stage2"
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():
    import httpx
    from sqlalchemy import select

    from app.config import Settings
    from app.main import create_app
    from app.models import ResearchExperiment, ResearchInputSnapshot, ResearchParent
    from tests.test_research_outcomes import (
        test_feature_conversion_keeps_original_version_through_experiment,
        test_lineage_retains_multiple_parents_and_descendants,
        test_saved_views_validate_and_retain_sort_columns,
        test_sync_and_model_advice_do_not_rewrite_report,
    )

    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:
            response = await client.post(
                "/api/v1/auth/login", json={"username": "admin", "password": "research-acceptance-only"}
            )
            assert response.status_code == 200
            async with app.state.sessions() as db:
                fixed = await db.scalar(select(ResearchInputSnapshot))
                for experiment in await db.scalars(select(ResearchExperiment)):
                    for parent in experiment.parents:
                        assert await db.get(ResearchParent, (experiment.id, parent["kind"], parent["id"]))
            await test_feature_conversion_keeps_original_version_through_experiment(client, {"id": fixed.id})
            await test_saved_views_validate_and_retain_sort_columns(client)
            await test_sync_and_model_advice_do_not_rewrite_report(app, client)
            await test_lineage_retains_multiple_parents_and_descendants(app, client, {"id": fixed.id})
    print("PASS PostgreSQL: feature revisions, saved views, immutable evaluations and multi-parent traversal")


if __name__ == "__main__":
    docker("createdb", "-U", "postgres", NAME)
    with Path("/tmp/wq-research-stage1.dump").open("rb") as source:
        docker("pg_restore", "-U", "postgres", "-d", NAME, stdin=source)
    config = Config("alembic.ini")
    command.upgrade(config, "0007")
    command.check(config)
    asyncio.run(acceptance())
    dump = Path("/tmp/wq-research-stage2.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_evaluations),(SELECT count(*) FROM research_parents),(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: 0006 → 0007 and pg_dump/pg_restore preserve history, graph and old research notes"
    )