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

"""Run against an isolated acceptance Compose project; never use a personal data stack.

From the repository root:
  python3 backend/tests/docker_acceptance.py --env-file .local/docker-test.env
Requires images already built and services started with -p wq-alpha-acceptance.
"""

import argparse
import http.cookiejar
import json
import os
import subprocess
import time
import urllib.error
import urllib.request
from pathlib import Path


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--env-file", required=True, type=Path)
    parser.add_argument("--project", default="wq-alpha-acceptance")
    args = parser.parse_args()
    if not args.project.startswith("wq-alpha-acceptance"):
        raise SystemExit("Only an isolated wq-alpha-acceptance* project is allowed")
    values = dict(
        line.split("=", 1)
        for line in args.env_file.read_text().splitlines()
        if "=" in line and not line.startswith("#")
    )
    compose = ["docker", "compose", "--env-file", str(args.env_file), "-p", args.project]

    def run(arguments, data=None):
        result = subprocess.run(compose + arguments, input=data, capture_output=True, check=True)
        return result.stdout

    # Assert that we are addressing the isolated project rather than an unrelated localhost app.
    published = run(["port", "web", "80"]).decode().strip()
    assert published == f"127.0.0.1:{values.get('LOCAL_PORT', '8080')}"
    base = f"http://localhost:{values.get('LOCAL_PORT', '8080')}"
    opener = urllib.request.build_opener(
        urllib.request.ProxyHandler({}), urllib.request.HTTPCookieProcessor(http.cookiejar.CookieJar())
    )

    def request(path, method="GET", payload=None):
        body = json.dumps(payload).encode() if payload is not None else None
        return opener.open(
            urllib.request.Request(
                base + path,
                data=body,
                method=method,
                headers={"Content-Type": "application/json", "X-WQ-Request": "1", "Origin": base},
            ),
            timeout=10,
        )

    assert json.load(request("/api/v1/health")) == {"status": "ok"}
    html = request("/")
    assert "frame-ancestors 'none'" in html.headers["Content-Security-Policy"]
    assert b'<div id="root">' in html.read()
    try:
        request("/api/v1/alphas")
        raise AssertionError("Unauthenticated data was exposed")
    except urllib.error.HTTPError as error:
        assert error.code == 401
    json.load(
        request(
            "/api/v1/auth/login",
            "POST",
            {"username": values.get("ADMIN_USERNAME", "admin"), "password": values["ADMIN_PASSWORD"]},
        )
    )
    code = b"""import asyncio
from app.config import Settings
from app.db import create_database
from app.alphas import upsert_alpha
async def seed():
    engine, sessions = create_database(Settings().database_url)
    async with sessions() as db:
        await upsert_alpha(db, {"id":"DOCKER_ACCEPTANCE", "name":"Synthetic acceptance record", "type":"REGULAR", "stage":"IS", "status":"UNSUBMITTED", "regular":{"code":"rank(close)"}, "settings":{"region":"USA", "language":"FASTEXPR"}, "is":{"sharpe":2.0}})
        await db.commit()
    await engine.dispose()
asyncio.run(seed())
"""
    run(["exec", "-T", "backend", "python", "-"], code)
    json.load(
        request(
            "/api/v1/alphas/DOCKER_ACCEPTANCE/research",
            "PATCH",
            {
                "note": "persistent local note",
                "tags": ["docker-verified"],
                "state": "candidate",
                "favorite": True,
            },
        )
    )
    print("PASS: PostgreSQL migration, login, API authorization, local research write")
    run(["up", "-d", "--force-recreate", "--wait"])
    # A persisted server session and all database rows survive container replacement.
    detail = json.load(request("/api/v1/alphas/DOCKER_ACCEPTANCE"))
    assert detail["research"]["note"] == "persistent local note"
    run(["exec", "-T", "backend", "python", "-"], code.replace(b'"sharpe":2.0', b'"sharpe":3.0'))
    detail = json.load(request("/api/v1/alphas/DOCKER_ACCEPTANCE"))
    assert detail["sharpe"] == 3 and detail["research"]["state"] == "candidate"
    assert detail["research"]["tags"] == ["docker-verified"]
    print("PASS: container replacement preserves session and data; snapshot update preserves research")
    dump = run(["exec", "-T", "db", "pg_dump", "-U", "wq", "-d", "wq", "-Fc", "--no-owner"])
    backup = Path(".local/docker-acceptance.dump")
    fd = os.open(backup, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "wb") as output:
        output.write(dump)
    # Only this script's scratch restore database is replaced on repeated acceptance runs.
    run(["exec", "-T", "db", "dropdb", "-U", "wq", "--if-exists", "wq_acceptance_restore"])
    run(["exec", "-T", "db", "createdb", "-U", "wq", "wq_acceptance_restore"])
    run(
        [
            "exec",
            "-T",
            "db",
            "pg_restore",
            "-U",
            "wq",
            "-d",
            "wq_acceptance_restore",
            "--no-owner",
            "--exit-on-error",
        ],
        dump,
    )
    rows = (
        run(
            [
                "exec",
                "-T",
                "db",
                "psql",
                "-U",
                "wq",
                "-d",
                "wq_acceptance_restore",
                "-At",
                "-c",
                "SELECT note || '|' || state FROM research WHERE alpha_id = 'DOCKER_ACCEPTANCE'",
            ]
        )
        .decode()
        .strip()
    )
    assert rows == "persistent local note|candidate"
    print("PASS: PostgreSQL custom-format backup restores records into independent database")
    config = json.loads(run(["-f", "compose.public.yaml", "config", "--format", "json"]))
    assert config["services"]["backend"]["environment"]["COOKIE_SECURE"] == "true"
    assert config["services"]["backend"]["environment"]["PUBLIC_ORIGIN"].startswith("https://")
    assert "ports" not in config["services"]["db"] and "ports" not in config["services"]["backend"]
    run(
        [
            "exec",
            "-T",
            "web",
            "caddy",
            "validate",
            "--config",
            "/etc/caddy/Caddyfile",
            "--adapter",
            "caddyfile",
        ]
    )
    print("PASS: Caddy configuration and public HTTPS/Secure-cookie configuration checks")
    # Keep the backend process alive while its database disappears and returns.
    run(["stop", "db"])
    run(["up", "-d", "--wait", "db"])
    probe = b"""import asyncio
from app.config import Settings
from app.db import create_database
from app.jobs import create_job
async def enqueue():
    engine, sessions = create_database(Settings().database_url)
    async with sessions() as db:
        job = await create_job(db, "profile")
        print(job.id)
    await engine.dispose()
asyncio.run(enqueue())
"""
    job_id = run(["exec", "-T", "backend", "python", "-"], probe).decode().strip()
    deadline = time.monotonic() + 15
    while time.monotonic() < deadline:
        status = json.load(request(f"/api/v1/sync-jobs/{job_id}"))["status"]
        if status == "waiting_connection":
            break
        time.sleep(0.25)
    else:
        raise AssertionError("Scheduler did not resume after database restart")
    print("PASS: live backend resumes task processing after database stop/start")


if __name__ == "__main__":
    main()