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",
                "version": 1,
                "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")

    # This synthetic service runs inside the disposable backend container. No real provider is used.
    model_server = Path(__file__).with_name("model_protocol.py").read_text()

    def start_model():
        run(["exec", "-d", "-T", "backend", "python", "-c", model_server])
        run(
            [
                "exec",
                "-T",
                "backend",
                "python",
                "-c",
                "import socket,time\nfor _ in range(40):\n try:\n  socket.create_connection(('127.0.0.1',19010),timeout=1).close();break\n except OSError: time.sleep(.1)\nelse: raise RuntimeError('Mock model did not start')",
            ]
        )

    start_model()
    for protocol in ("chat_completions", "responses"):
        config = {"base_url": "http://127.0.0.1:19010/v1", "model": "mock-model", "protocol": protocol}
        json.load(request("/api/v1/ai/settings", "PUT", config | {"api_key": "synthetic-model-key"}))
        tested = json.load(request("/api/v1/ai/settings/test", "POST"))
        assert tested["ready"], tested
    json.load(request("/api/v1/ai/settings", "PUT", config | {"enabled": True}))
    conversation = json.load(request("/api/v1/ai/conversations", "POST"))["id"]
    started = time.monotonic()
    stream = request(
        f"/api/v1/ai/conversations/{conversation}/runs", "POST", {"request_id": "stream", "message": "SLOW"}
    )
    run_id = stream.headers["X-AI-Run-ID"]
    assert stream.headers["x-vercel-ai-ui-message-stream"] == "v1"
    while b'"text-delta"' not in stream.readline():
        assert time.monotonic() - started < 8
    assert json.load(request(f"/api/v1/ai/runs/{run_id}"))["status"] == "running"
    stream.close()
    for _ in range(40):
        if json.load(request(f"/api/v1/ai/runs/{run_id}"))["status"] == "completed":
            break
        time.sleep(0.25)
    else:
        raise AssertionError("Disconnected execution did not complete")
    stream = request(
        f"/api/v1/ai/conversations/{conversation}/runs",
        "POST",
        {"request_id": "write", "message": "修改研究记录"},
    )
    run_id = stream.headers["X-AI-Run-ID"]
    stream.read()
    pending = json.load(request(f"/api/v1/ai/runs/{run_id}"))
    assert pending["status"] == "waiting_approval", pending
    assert (
        json.load(request("/api/v1/alphas/DOCKER_ACCEPTANCE"))["research"]["note"] == "persistent local note"
    )
    run(["up", "-d", "--force-recreate", "--wait"])
    start_model()
    assert json.load(request("/api/v1/ai/settings"))["ready"]
    approval = pending["tools"][0]["id"]
    for _ in range(2):
        request(f"/api/v1/ai/approvals/{approval}/decision", "POST", {"approved": True}).read()
    saved = json.load(request("/api/v1/alphas/DOCKER_ACCEPTANCE"))["research"]
    assert saved["note"] == "AI verified note" and saved["version"] == 3
    try:
        request("/api/v1/alphas/DOCKER_ACCEPTANCE/research", "PATCH", {"version": 2, "note": "stale"})
        raise AssertionError("A stale edit overwrote AI research")
    except urllib.error.HTTPError as error:
        assert error.code == 409
    print(
        "PASS: both real SDK protocols over mock HTTP; Caddy streams before completion; disconnect recovery"
    )
    print(
        "PASS: pending approval survives container replacement; duplicate confirmation writes once; PostgreSQL version conflict"
    )
    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 == "AI verified note|candidate"
    ai_rows = (
        run(
            [
                "exec",
                "-T",
                "db",
                "psql",
                "-U",
                "wq",
                "-d",
                "wq_acceptance_restore",
                "-At",
                "-c",
                "SELECT (SELECT count(*) FROM ai_conversations), (SELECT count(*) FROM ai_runs), (SELECT count(*) FROM ai_tool_calls), (SELECT count(*) FROM ai_settings WHERE api_key_encrypted IS NOT NULL)",
            ]
        )
        .decode()
        .strip()
    )
    assert ai_rows == "1|2|1|1", ai_rows
    print("PASS: PostgreSQL custom-format backup restores records into independent database")
    # Downgrade only the scratch restore database, then verify upgrading existing 0001 research.
    migration_env = (
        "DATABASE_URL=postgresql+asyncpg://wq:"
        + values["POSTGRES_PASSWORD"]
        + "@db:5432/wq_acceptance_restore"
    )
    for command in (["downgrade", "0001"], ["upgrade", "head"], ["check"]):
        run(["exec", "-T", "-e", migration_env, "backend", "alembic", *command])
    versioned = (
        run(
            [
                "exec",
                "-T",
                "db",
                "psql",
                "-U",
                "wq",
                "-d",
                "wq_acceptance_restore",
                "-At",
                "-c",
                "SELECT note || '|' || version FROM research WHERE alpha_id='DOCKER_ACCEPTANCE'",
            ]
        )
        .decode()
        .strip()
    )
    assert versioned == "AI verified note|1", versioned
    print("PASS: existing 0001 research upgrades with version 1 and unchanged content; Alembic model parity")
    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()