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

"""One-off acceptance against the dedicated local PostgreSQL catalog_test database."""

import asyncio
import os
import re

from alembic import command
from alembic.config import Config
from cryptography.fernet import Fernet
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine

database_name = os.environ.get("WQ_CATALOG_ACCEPTANCE_DATABASE", "catalog_flow_test")
if not re.fullmatch(r"catalog_[a-z0-9_]{1,40}", database_name):
    raise ValueError("Acceptance requires a dedicated catalog_* database")
URL = f"postgresql+asyncpg://postgres:catalog-test-only@127.0.0.1:18436/{database_name}"
os.environ.update(
    DATABASE_URL=URL, ADMIN_PASSWORD="migration-test-only", ENCRYPTION_KEY=Fernet.generate_key().decode()
)


async def sql(statement):
    engine = create_async_engine(URL)
    async with engine.begin() as connection:
        result = await connection.execute(text(statement))
        value = result.fetchall() if result.returns_rows else None
    await engine.dispose()
    return value


if __name__ == "__main__":
    config = Config("alembic.ini")
    if asyncio.run(sql("SELECT tablename FROM pg_tables WHERE schemaname='public'")):
        raise RuntimeError("Acceptance database must be empty; existing data will not be overwritten")
    command.upgrade(config, "0002")
    asyncio.run(
        sql(
            "INSERT INTO alphas (id, hidden, settings, is_metrics, os_metrics, checks, synced_at, raw) VALUES ('MIGRATION_TEST', false, '{}', '{}', '{}', '[]', now(), '{}');"
        )
    )
    asyncio.run(
        sql(
            "INSERT INTO research (alpha_id, note, tags, favorite, state, updated_at, version) VALUES ('MIGRATION_TEST', 'preserve research', '[]', false, 'inbox', now(), 7);"
        )
    )
    command.upgrade(config, "head")
    command.check(config)
    assert asyncio.run(sql("SELECT note, version FROM research WHERE alpha_id='MIGRATION_TEST'")) == [
        ("preserve research", 7)
    ]
    assert asyncio.run(sql("SELECT count(*) FROM catalog_batches")) == [(0,)]
    command.downgrade(config, "0002")
    command.upgrade(config, "head")
    command.check(config)
    assert asyncio.run(sql("SELECT note, version FROM research WHERE alpha_id='MIGRATION_TEST'")) == [
        ("preserve research", 7)
    ]
    print(
        "PostgreSQL 17: 0002 → 0003, downgrade/re-upgrade, metadata check, Alpha/research preservation passed"
    )

    async def flow():
        import httpx

        from app.config import Settings
        from app.main import create_app
        from app.worldquant import WqClient
        from tests.catalog_fake import catalog_response
        from tests.test_catalog import SCOPE, prepare, search, sync

        def upstream(request):
            if request.url.path == "/authentication":
                return httpx.Response(201, json={"token": {"expiry": 14400}})
            if request.url.path == "/users/self":
                return httpx.Response(200, json={"id": "PG_TEST_USER"})
            assert request.method == "GET"
            return catalog_response(request) or httpx.Response(404)

        settings = Settings(_env_file=None, enable_runner=False, public_origin="http://testserver")
        app = create_app(settings, WqClient(settings, transport=httpx.MockTransport(upstream)))
        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": "migration-test-only"}
                    )
                ).status_code == 200
                await client.put(
                    "/api/v1/account/credentials",
                    json={"email": "pg@example.com", "password": "synthetic-only"},
                )
                job = (await client.post("/api/v1/account/connect")).json()
                await app.state.runner.execute(job["id"])
                catalog = (client, app.state.runner, {})
                assert (await sync(catalog))["status"] == "completed"
                version = (await sync(catalog, "TEST_FIN"))["id"]
                result = await search(client, "/datasets/TEST_FIN/fields")
                assert result["complete_count"] == 123
                draft = (await prepare(client, version)).json()
                assert len(draft["field_ids"]) == 123
                responses = await asyncio.gather(
                    *[
                        client.patch(
                            "/api/v1/catalog/datasets/TEST_FIN/research",
                            params=SCOPE,
                            json={"version": 1, "note": value},
                        )
                        for value in ["one", "two"]
                    ]
                )
                assert sorted(r.status_code for r in responses) == [200, 409]
                await sync(catalog, "TEST_FIN")
                assert (await prepare(client, version)).status_code == 409
                persisted = (await client.get("/api/v1/catalog/inputs/" + draft["id"])).json()
                assert persisted == draft
        print(
            "PostgreSQL: real API/runner multi-page dedupe, immutable draft, refresh conflict and concurrent note CAS passed"
        )

    asyncio.run(flow())