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_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()