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

import csv
import io
from datetime import timedelta

import pytest
from sqlalchemy import select

from app.alphas import upsert_alpha
from app.models import Account, Admin, LoginSession, Research, now
from app.security import cipher
from tests.conftest import ADMIN_PASSWORD, alpha

PREFIX = "/api/v1"


async def seed(app, count=1):
    async with app.state.sessions() as db:
        for i in range(count):
            await upsert_alpha(db, alpha(f"a{i:04}", name=f"Signal {i:04}", **{"is": {"sharpe": i / 100}}))
        await db.commit()


async def test_login_logout_expiry_and_password_hash(app, client):
    assert (await client.get(f"{PREFIX}/alphas")).status_code == 401
    wrong = await client.post(f"{PREFIX}/auth/login", json={"username": "admin", "password": "wrong"})
    assert wrong.status_code == 401
    login = await client.post(f"{PREFIX}/auth/login", json={"username": "admin", "password": ADMIN_PASSWORD})
    assert login.status_code == 200
    assert "HttpOnly" in login.headers["set-cookie"] and "SameSite=strict" in login.headers["set-cookie"]
    token = client.cookies.get("wq_session")
    async with app.state.sessions() as db:
        admin = await db.get(Admin, 1)
        session = (await db.scalars(select(LoginSession))).one()
        assert admin.password_hash.startswith("$argon2id$") and ADMIN_PASSWORD not in admin.password_hash
        assert session.token_hash != token
    assert (await client.post(f"{PREFIX}/auth/logout")).status_code == 200
    client.cookies.set("wq_session", token)
    assert (await client.get(f"{PREFIX}/auth/me")).status_code == 401
    await client.post(f"{PREFIX}/auth/login", json={"username": "admin", "password": ADMIN_PASSWORD})
    async with app.state.sessions() as db:
        session = (await db.scalars(select(LoginSession))).one()
        session.expires_at = now() - timedelta(seconds=1)
        await db.commit()
    assert (await client.get(f"{PREFIX}/auth/me")).status_code == 401


async def test_login_rate_limit_and_csrf(client):
    for _ in range(5):
        assert (
            await client.post(f"{PREFIX}/auth/login", json={"username": "admin", "password": "wrong"})
        ).status_code == 401
    assert (
        await client.post(f"{PREFIX}/auth/login", json={"username": "admin", "password": ADMIN_PASSWORD})
    ).status_code == 429
    assert (
        await client.post(f"{PREFIX}/auth/login", headers={"X-WQ-Request": ""}, json={})
    ).status_code == 403
    assert (
        await client.post(f"{PREFIX}/auth/login", headers={"Origin": "https://evil.example"}, json={})
    ).status_code == 403


async def test_public_session_cookie_is_secure(app, client):
    app.state.settings.cookie_secure = True
    response = await client.post(
        f"{PREFIX}/auth/login", json={"username": "admin", "password": ADMIN_PASSWORD}
    )
    assert response.status_code == 200
    assert "; Secure" in response.headers["set-cookie"] and "; HttpOnly" in response.headers["set-cookie"]


async def test_credentials_encrypted_and_never_returned(app, logged_in):
    secret = "private-worldquant-password"
    r = await logged_in.put(
        f"{PREFIX}/account/credentials", json={"email": "person@example.com", "password": secret}
    )
    assert r.status_code == 200 and r.json()["configured"] is True
    assert "password" not in r.text and secret not in r.text
    async with app.state.sessions() as db:
        account = await db.get(Account, 1)
        assert secret not in account.password_encrypted
        assert cipher(app.state.settings).decrypt(account.password_encrypted.encode()).decode() == secret
        account.wq_user_id = "person"
        await db.commit()
    assert (
        await logged_in.put(
            f"{PREFIX}/account/credentials", json={"email": "other@example.com", "password": secret}
        )
    ).status_code == 409
    invalid = await logged_in.put(
        f"{PREFIX}/account/credentials", json={"email": "invalid", "password": secret}
    )
    assert invalid.status_code == 422 and secret not in invalid.text
    response = await logged_in.get(f"{PREFIX}/account")
    assert response.json()["profile"] == {} and "password" not in response.text


async def test_preferences_and_openapi(app, logged_in):
    values = {"display_name": "My Lab", "theme": "dark", "timezone": "UTC", "page_size": 100}
    r = await logged_in.patch(f"{PREFIX}/account/preferences", json=values)
    assert r.status_code == 200 and all(r.json()[k] == v for k, v in values.items())
    values["timezone"] = "invalid/timezone"
    assert (await logged_in.patch(f"{PREFIX}/account/preferences", json=values)).status_code == 422
    contract = app.openapi()
    assert contract["paths"]["/api/v1/alphas"]["get"]["security"] == [{"APIKeyCookie": []}]
    assert "text/csv" in contract["paths"]["/api/v1/alphas/export"]["get"]["responses"]["200"]["content"]
    assert contract["paths"]["/api/v1/alphas"]["get"]["responses"]["422"]["content"]["application/json"][
        "schema"
    ]["$ref"].endswith("/ErrorOutput")


async def test_local_research_survives_snapshot_and_bulk_is_atomic(app, logged_in):
    await seed(app, 3)
    r = await logged_in.patch(
        f"{PREFIX}/alphas/a0000/research",
        json={"version": 1, "note": "keep hypothesis", "tags": ["a", "a"], "state": "candidate", "favorite": True},
    )
    assert r.status_code == 200
    async with app.state.sessions() as db:
        await upsert_alpha(db, alpha("a0000", **{"is": {"sharpe": 9.25}}))
        await db.commit()
    detail = (await logged_in.get(f"{PREFIX}/alphas/a0000")).json()
    assert detail["sharpe"] == 9.25 and detail["fitness"] is None
    assert detail["research"]["note"] == "keep hypothesis" and detail["research"]["favorite"]
    assert detail["research"]["tags"] == ["a"] and detail["research"]["state"] == "candidate"
    invalid = await logged_in.patch(
        f"{PREFIX}/alphas/research/bulk", json={"alpha_ids": ["a0000", "missing"], "add_tags": ["bad"], "versions": {"a0000": 2, "missing": 1}}
    )
    assert invalid.status_code == 404
    assert (await logged_in.get(f"{PREFIX}/alphas?tag=bad")).json()["total"] == 0
    r = await logged_in.patch(
        f"{PREFIX}/alphas/research/bulk",
        json={
            "alpha_ids": ["a0000", "a0001"],
            "versions": {"a0000": 2, "a0001": 1},
            "add_tags": ["new"],
            "remove_tags": ["a"],
            "state": "optimizing",
        },
    )
    assert r.json()["updated"] == 2
    result = (await logged_in.get(f"{PREFIX}/alphas?tag=new&research_state=optimizing")).json()
    assert result["total"] == 2
    assert (await logged_in.get(f"{PREFIX}/alphas?tag=ne")).json()["total"] == 0
    await logged_in.patch(f"{PREFIX}/alphas/a0000/research", json={"note": "updated", "version": 3})
    detail = (await logged_in.get(f"{PREFIX}/alphas/a0000")).json()
    assert detail["research"]["favorite"] and detail["research"]["tags"] == ["new"]


async def test_filter_pagination_sort_and_export_over_500(app, logged_in):
    await seed(app, 603)
    query = "region=USA&sharpe_min=0.02&sharpe_max=6&sort=sharpe&direction=desc"
    page = (await logged_in.get(f"{PREFIX}/alphas?{query}&offset=10&limit=25")).json()
    assert page["total"] == 599 and len(page["items"]) == 25
    assert page["items"][0]["id"] == "a0590"
    exported = await logged_in.get(f"{PREFIX}/alphas/export?{query}&offset=10&limit=25")
    rows = list(csv.DictReader(io.StringIO(exported.text.lstrip("\ufeff"))))
    assert len(rows) == 599 and rows[0]["id"] == "a0600" and rows[-1]["id"] == "a0002"
    assert (await logged_in.get(f"{PREFIX}/alphas?q=Signal%200001")).json()["total"] == 1
    assert (await logged_in.get(f"{PREFIX}/alphas?q=%25")).json()["total"] == 0
    for query in (
        "sort=raw",
        "sharpe_min=nan",
        "sharpe_min=inf",
        "sharpe_min=5&sharpe_max=2",
        "limit=1000",
        "created_from=2026-01-01&created_to=2025-01-01T00:00:00Z",
    ):
        assert (await logged_in.get(f"{PREFIX}/alphas?{query}")).status_code == 422


async def test_export_formula_injection_and_detail_variants(app, logged_in):
    async with app.state.sessions() as db:
        await upsert_alpha(
            db,
            alpha(
                "super1",
                name="=DANGEROUS()",
                type="SUPER",
                regular=None,
                selection={"code": "self_correlation < 0.5"},
                combo={"code": "alpha"},
                settings={"language": "FASTEXPR"},
                **{"is": {"sharpe": None, "fitness": float("nan")}},
            ),
        )
        await upsert_alpha(
            db,
            alpha(
                "python1", regular={"code": "def alpha():\n    return close"}, settings={"language": "PYTHON"}
            ),
        )
        await db.commit()
        row = await db.get(Research, "super1")
        row.note = " @formula()"
        await db.commit()
    detail = (await logged_in.get(f"{PREFIX}/alphas/super1")).json()
    assert detail["selection"] == "self_correlation < 0.5" and detail["combo"] == "alpha"
    assert detail["expression"] is None and detail["sharpe"] is None and detail["fitness"] is None
    assert (await logged_in.get(f"{PREFIX}/alphas/python1")).json()["language"] == "PYTHON"
    csv_data = (await logged_in.get(f"{PREFIX}/alphas/export?q=super1")).text
    row = list(csv.DictReader(io.StringIO(csv_data.lstrip("\ufeff"))))[0]
    assert row["name"] == "'=DANGEROUS()" and row["note"] == "' @formula()"
    assert (await logged_in.get(f"{PREFIX}/alphas/super1/pnl")).json() == {
        "cached": False,
        "series": [],
        "points": [],
        "fetched_at": None,
    }


async def test_job_api_deduplicates_and_retries(app, logged_in):
    assert (await logged_in.post(f"{PREFIX}/sync-jobs", json={"kind": "full_sync"})).status_code == 409
    await logged_in.put(
        f"{PREFIX}/account/credentials", json={"email": "person@example.com", "password": "only-testing"}
    )
    connect = await logged_in.post(f"{PREFIX}/account/connect")
    assert connect.status_code == 202
    first = (await logged_in.post(f"{PREFIX}/sync-jobs", json={"kind": "full_sync"})).json()
    again = (await logged_in.post(f"{PREFIX}/sync-jobs", json={"kind": "full_sync"})).json()
    assert first["id"] == again["id"]
    await logged_in.post(f"{PREFIX}/sync-jobs/{first['id']}/cancel")
    assert (await logged_in.get(f"{PREFIX}/sync-jobs/{first['id']}")).json()["status"] == "cancelled"
    assert (await logged_in.post(f"{PREFIX}/sync-jobs/{first['id']}/retry")).json()["status"] == "queued"
    assert (
        await logged_in.post(
            f"{PREFIX}/sync-jobs", json={"kind": "alpha_refresh", "alpha_ids": ["../authentication"]}
        )
    ).status_code == 422


@pytest.mark.parametrize("endpoint", ["account", "alphas", "alphas/export", "alphas/a1/pnl", "sync-jobs"])
async def test_every_data_endpoint_requires_session(client, endpoint):
    assert (await client.get(f"{PREFIX}/{endpoint}")).status_code == 401