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

81 lines
3.7 KiB
Python
Raw Normal View History

"""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, ResearchParent, TemplateInput
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(TemplateInput))
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"
)