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

import asyncio
from datetime import datetime, timedelta, timezone
from email.utils import format_datetime

import httpx
import pytest

from app.alphas import pnl_points, sanitize
from app.worldquant import VerificationRequired, WqClient, WqError


async def test_cookie_auth_expiry_and_read_only_boundary(settings):
    calls, auth_count = [], 0

    def handler(request):
        nonlocal auth_count
        calls.append((request.method, request.url.path))
        if request.url.path == "/authentication":
            auth_count += 1
            return httpx.Response(201, headers={"Set-Cookie": f"t=session{auth_count}; Path=/"}, json={})
        assert "Authorization" not in request.headers
        if request.headers.get("Cookie") == "t=session1":
            return httpx.Response(401)
        return httpx.Response(200, json={"id": "user1"})

    client = WqClient(settings, transport=httpx.MockTransport(handler))
    await client.authenticate("person@example.com", "secret")
    assert await client.profile() == {"id": "user1"}
    assert auth_count == 2
    await asyncio.gather(*(client.authenticate("person@example.com", "secret") for _ in range(5)))
    assert auth_count == 2
    assert all(method == "GET" or path == "/authentication" for method, path in calls)
    client.disconnect()
    assert client.credentials is None and not list(client.client.cookies.items())
    await client.close()


async def test_auth_metadata_usage_and_local_simulation_allowance(settings):
    from datetime import datetime
    from zoneinfo import ZoneInfo

    today = datetime.now(ZoneInfo("America/New_York")).date().isoformat()
    calls = []

    def handler(request):
        calls.append((request.method, request.url.path))
        if request.url.path == "/authentication":
            return httpx.Response(
                201,
                json={
                    "permissions": ["CONSULTANT", "SUPER_ALPHA", "CONSULTANT"],
                    "token": {"expiry": 14400, "value": "private-token"},
                    "password": "private-password",
                },
            )
        if request.url.path.endswith("/simulations"):
            return httpx.Response(
                200,
                json={
                    "records": {
                        "schema": {"properties": [{"name": "value"}, {"name": "date"}]},
                        "records": [[0, today]],
                    },
                    "yesterday": {"value": 40},
                    "total": {"value": 900},
                },
            )
        if request.url.path.endswith("/submissions"):
            return httpx.Response(403)
        assert request.headers["Accept"] == "application/json;version=4.0"
        return httpx.Response(200, json={"active": 99, "unsubmitted": 100, "decommissioned": 0})

    client = WqClient(settings, transport=httpx.MockTransport(handler))
    await client.authenticate("test@example.com", "secret")
    assert client.permissions == ["CONSULTANT", "SUPER_ALPHA"]
    info = client.session_info()
    assert info["authenticated"] and 14395 < info["remaining_seconds"] <= 14400
    assert "private" not in str(info)
    usage = await client.account_usage()
    assert usage["simulations"]["today"] == 0
    assert usage["simulations"]["limit"] == 10_000 and usage["simulations"]["remaining"] == 10_000
    assert usage["alphas"]["active"] == 99
    assert "submissions" in usage["errors"] and "submissions" not in usage
    assert all(method == "GET" or path == "/authentication" for method, path in calls)
    client.disconnect()
    assert client.permissions is None and client.session_info()["remaining_seconds"] is None
    assert not client.session_info()["authenticated"]
    await client.close()


def test_activity_missing_today_is_unknown_not_zero():
    from app.account_data import daily_activity

    data = {
        "records": {
            "schema": {"properties": [{"name": "date"}, {"name": "value"}]},
            "records": [["2026-09-06", 2]],
        }
    }
    assert daily_activity(data, "2026-09-07")["today"] is None
    assert daily_activity({}, "2026-09-07")["limit"] is None
    local_budget = daily_activity(data, "2026-09-07", daily_limit=10_000)
    assert local_budget["limit"] == 10_000 and local_budget["remaining"] is None


@pytest.mark.parametrize(
    "recordset", [None, {"records": None}, {"schema": None}, {"schema": {"properties": None}}]
)
def test_malformed_optional_activity_does_not_discard_other_sections(recordset):
    from app.account_data import usage_snapshot

    result = usage_snapshot(
        {"simulations": {"records": recordset}, "submissions": {}, "alphas": {"active": 99}}, {}
    )
    assert "simulations" in result["errors"] and "simulations" not in result
    assert result["submissions"]["today"] is None and result["alphas"]["active"] == 99


async def test_persona_verification_uses_same_cookie_and_safe_location(settings):
    complete = False

    def handler(request):
        if request.url.path == "/authentication":
            return httpx.Response(
                401,
                headers={
                    "WWW-Authenticate": "persona",
                    "Location": "/authentication/persona/challenge",
                    "Set-Cookie": "challenge=abc; Path=/",
                },
            )
        assert request.headers.get("Cookie") == "challenge=abc"
        return httpx.Response(201 if complete else 202, json={})

    client = WqClient(settings, transport=httpx.MockTransport(handler))
    with pytest.raises(VerificationRequired) as error:
        await client.authenticate("test@example.com", "secret")
    assert error.value.url.endswith("/authentication/persona/challenge")
    with pytest.raises(VerificationRequired):
        await client.verify()
    complete = True
    await client.verify()
    assert client.authenticated and client.verification_url is None
    await client.close()
    unsafe = WqClient(
        settings,
        transport=httpx.MockTransport(
            lambda r: httpx.Response(
                401, headers={"WWW-Authenticate": "persona", "Location": "https://evil.example/secret"}
            )
        ),
    )
    with pytest.raises(WqError) as error:
        await unsafe.authenticate("test@example.com", "secret")
    assert error.value.code == "invalid_verification"
    await unsafe.close()


async def test_retry_after_network_budget_pending_and_permissions(settings):
    delays, requests = [], 0

    async def sleep(delay):
        delays.append(delay)

    def handler(request):
        nonlocal requests
        requests += 1
        if requests == 1:
            return httpx.Response(429, headers={"Retry-After": "120"})
        if requests == 2:
            raise httpx.ConnectError("contains private body", request=request)
        return httpx.Response(200, json={"id": "ready"})

    client = WqClient(settings, transport=httpx.MockTransport(handler), sleep=sleep)
    client.credentials, client.authenticated = ("test@example.com", "secret"), True
    assert await client.profile() == {"id": "ready"}
    assert delays[0] == 120 and len(delays) == 2
    date = format_datetime(datetime.now(timezone.utc) + timedelta(seconds=60), usegmt=True)
    assert 58 < client.retry_delay(date, 0) <= 60
    await client.close()
    for status, code in [(403, "access_denied"), (404, "not_found"), (503, "retry_exhausted")]:
        c = WqClient(
            settings,
            transport=httpx.MockTransport(lambda r: httpx.Response(status, text="private data")),
            sleep=sleep,
        )
        c.credentials, c.authenticated = ("test@example.com", "secret"), True
        with pytest.raises(WqError) as error:
            await c.profile()
        assert error.value.code == code and "private" not in str(error.value)
        await c.close()
    responses = iter(
        [httpx.Response(202, headers={"Retry-After": "1"}), httpx.Response(200, json={"records": []})]
    )
    c = WqClient(settings, transport=httpx.MockTransport(lambda r: next(responses)), sleep=sleep)
    c.credentials, c.authenticated = ("test@example.com", "secret"), True
    assert await c.pnl("abc") == {"records": []}
    await c.close()


@pytest.mark.parametrize(
    "raw,expected",
    [
        (
            {
                "schema": {"properties": [{"name": "date"}, {"name": "pnl"}]},
                "records": [["2025-01-01", 10], ["2025-01-02", None]],
            },
            [{"date": "2025-01-01", "value": 10.0}, {"date": "2025-01-02", "value": None}],
        ),
        (
            {"schema": {"properties": [{"name": "pnl"}, {"name": "date"}]}, "records": [[20, "2025-01-01"]]},
            [{"date": "2025-01-01", "value": 20.0}],
        ),
        (
            {"schema": {"properties": {"date": {}, "pnl": {}}}, "records": [["2025-01-01", "NaN"]]},
            [{"date": "2025-01-01", "value": None}],
        ),
        (
            {"records": [{"timestamp": 1735689600000, "value": "5"}]},
            [{"date": "2025-01-01T00:00:00+00:00", "value": 5.0}],
        ),
    ],
)
def test_pnl_schemas_preserve_null(raw, expected):
    assert pnl_points(raw) == expected


def test_unknown_pnl_schema_fails_instead_of_fabricating():
    with pytest.raises(ValueError):
        pnl_points({"records": [["2025-01-01", 100]]})


def test_sensitive_response_keys_are_removed_recursively():
    assert sanitize(
        {
            "id": "a1",
            "accessToken": "secret",
            "nested": [{"refresh_token": "secret", "client-secret": "secret", "value": 1}],
            "Set-Cookie": "secret",
        }
    ) == {"id": "a1", "nested": [{"value": 1}]}


async def test_concurrent_expiry_authenticates_once(settings):
    auth_count = 0
    first_requests = 0
    both_expired = asyncio.Event()

    async def handler(request):
        nonlocal auth_count, first_requests
        if request.method == "POST":
            auth_count += 1
            return httpx.Response(201, json={})
        if first_requests < 2:
            first_requests += 1
            if first_requests == 2:
                both_expired.set()
            await both_expired.wait()
            return httpx.Response(401)
        return httpx.Response(200, json={"id": "user"})

    client = WqClient(settings, transport=httpx.MockTransport(handler))
    await client.authenticate("test@example.com", "secret")
    assert await asyncio.gather(client.profile(), client.profile()) == [{"id": "user"}, {"id": "user"}]
    assert auth_count == 2  # Initial connection plus one shared reauthentication.
    await client.close()


async def test_slow_pnl_waits_beyond_network_retry_budget(settings):
    calls, delays = 0, []

    async def sleep(delay):
        delays.append(delay)

    def handler(request):
        nonlocal calls
        calls += 1
        assert request.url.path == '/alphas/LL977PqL/recordsets/pnl'
        if calls <= 6:
            return httpx.Response(200, headers={'Retry-After': '1.0'})
        return httpx.Response(200, json={'records': []})

    client = WqClient(settings, transport=httpx.MockTransport(handler), sleep=sleep)
    client.credentials, client.authenticated = ('test@example.com', 'secret'), True
    try:
        assert await client.pnl('LL977PqL') == {'records': []}
        assert calls == 7
        assert delays == [1.0] * 6
    finally:
        await client.close()


async def test_pnl_polling_remains_bounded(settings):
    settings.pnl_poll_attempts = 5
    delays = []

    async def sleep(delay):
        delays.append(delay)

    client = WqClient(
        settings,
        transport=httpx.MockTransport(lambda r: httpx.Response(200, headers={'Retry-After': '1.0'})),
        sleep=sleep,
    )
    client.credentials, client.authenticated = ('test@example.com', 'secret'), True
    try:
        with pytest.raises(WqError) as error:
            await client.pnl('LL977PqL')
        assert error.value.code == 'pending'
        assert delays == [1.0] * 4
    finally:
        await client.close()