tests/test_api.py

import asyncio
import base64
import io
import struct
import time
from types import SimpleNamespace

import httpx
import pytest
from PIL import Image
from pydantic import SecretStr, ValidationError

from backend.main import create_app
from backend.schemas import AVChunk
from backend.settings import Settings, validate_service_url


def start_body(duration=1000):
    return {"holdings": [{"id": "xom", "name": "Exxon Mobil Corporation", "ticker": "XOM",
                          "exchange": "NYSE", "currency": "USD", "asset_type": "equity", "weight": 100}],
            "objective": "Research how the reported oil supply conditions might affect this holding.",
            "duration_ms": duration, "source_name": "authorized-fixture.mp4", "authorized": True}


def media_chunk(sequence=0, count=16000, epoch=0):
    image = io.BytesIO()
    Image.new("RGB", (64, 36), (100 + sequence % 100, 100, 100)).save(image, "JPEG")
    # Non-silent alternating signal goes through base64/float32 validation.
    audio = struct.pack("<" + "f" * count, *[0.25 if i % 2 else -0.25 for i in range(count)])
    begin = sequence * 1000
    return {"epoch": epoch, "sequence": sequence, "source_start_ms": begin,
            "source_end_ms": begin + count / 16, "sample_rate": 16000, "sample_count": count,
            "audio_b64": base64.b64encode(audio).decode(), "final_chunk": count < 16000,
            "capture_time_ms": 1750000000000.125,
            "frames": [{"source_time_ms": begin + 0.0625, "jpeg_b64": base64.b64encode(image.getvalue()).decode()}]}


@pytest.mark.asyncio
async def test_connected_api_upload_partial_completion_and_retained_results():
    app = create_app(Settings(mock_latency_seconds=0))
    async with app.router.lifespan_context(app):
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app, client=("127.0.0.1", 44)), base_url="http://test") as client:
            result = await client.post("/api/sessions", json=start_body(1062.5))
            assert result.status_code == 200, result.text
            credentials = result.json()
            base = "/api/sessions/" + credentials["session_id"]
            headers = {"X-Session-Token": credentials["session_token"]}
            assert (await client.get(base)).status_code == 404
            assert (await client.get(base + '/timings')).status_code == 404
            assert (await client.get(base + '?after_revision=0&wait_seconds=20')).status_code == 404
            for payload in [media_chunk(), media_chunk(1, 1000)]:
                response = await client.post(base + "/chunks", headers=headers, json=payload)
                assert response.status_code == 200, response.text
                repeated = await client.post(base + "/chunks", headers=headers, json=payload)
                assert repeated.json()["duplicate"]
            assert (await client.post(base + "/complete", headers=headers)).status_code == 200
            session = app.state.sessions[credentials["session_id"]]
            await session.completion_task
            live_response = await client.get(base + f'?after_revision={session.revision}&wait_seconds=20', headers=headers)
            assert live_response.status_code == 200  # Terminal state returns without waiting.
            assert live_response.headers['cache-control'] == 'no-store'
            assert credentials['session_token'] not in live_response.text
            assert (await client.get(base + '?after_revision=0&wait_seconds=99', headers=headers)).status_code == 422
            snapshot = (await client.get(base, headers=headers)).json()
            assert snapshot["state"] == "completed"
            assert snapshot["coverage"]["processed_chunks"] == 2
            assert snapshot["coverage"]["processed_audio_ms"] == 1062.5
            assert snapshot["coverage"]["complete"]
            assert snapshot["evidence"][0]["source_start_ms"] == 0
            assert snapshot["assessments"][0]["implications"][0]["holding_id"] == "xom"
            report = await client.get(base + '/timings', headers=headers)
            assert report.status_code == 200 and report.headers['cache-control'] == 'no-store'
            assert len(report.json()['units']) == 2
            assert report.json()['summary']['contiguous_processed_source_ms'] == 1062.5
            for private in (credentials['session_token'], 'Exxon', 'authorized-fixture.mp4', 'audio_b64', 'jpeg_b64'):
                assert private not in report.text
            assert "discussion" not in snapshot
            assert snapshot["coverage"]["final_summary_requests"] == 0
            assert snapshot["brief"]["simulated"]
            assert snapshot["brief"]["complete"]
            assert credentials["session_token"] not in str(snapshot)
            assert (await client.post(base + "/stop", headers=headers)).json()["brief"] == snapshot["brief"]
            diagnostic = {"epoch": 1, "code": "audio_capture_gap", "message": "fixture", "source_start_ms": 0, "source_end_ms": 8}
            assert (await client.post(base + "/diagnostics", headers=headers, json=diagnostic)).status_code == 409
            diagnostic["epoch"] = 0
            assert (await client.post(base + "/diagnostics", headers=headers, json=diagnostic)).status_code == 200
            updated = (await client.get(base, headers=headers)).json()
            assert not updated["brief"]["complete"]
            assert updated["coverage"]["gaps"][-1]["source_end_ms"] == 8
            assert updated["coverage"]["final_summary_requests"] == 0



@pytest.mark.asyncio
async def test_auth_origin_url_and_body_safeguards():
    cfg = Settings(app_access_password=SecretStr("test-workspace-password-only"))
    app = create_app(cfg)
    async with app.router.lifespan_context(app):
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client:
            assert (await client.post("/api/sessions", json=start_body())).status_code == 401
            response = await client.post("/api/login", json={"password": "test-workspace-password-only"})
            assert response.status_code == 200, response.text
            assert "HttpOnly" in response.headers["set-cookie"]
            assert "test-workspace-password-only" not in response.text
            forbidden = await client.post("/api/sessions", json=start_body(), headers={"origin": "https://hostile.example"})
            assert forbidden.status_code == 403
            too_large = await client.post("/api/sessions", content=b"x" * 1_650_001)
            assert too_large.status_code == 413
            result = await client.post("/api/sessions", json=start_body())
            assert result.status_code == 200, result.text
    for url in ["http://169.254.169.254", "https://example.com", "https://foo.hf.space/path", "https://x:y@foo.hf.space", "https://foo.hf.space?token=bad"]:
        with pytest.raises(ValueError):
            validate_service_url(url, True)
    assert validate_service_url("https://test.hf.space") == "https://test.hf.space"


@pytest.mark.asyncio
async def test_workspace_login_starts_without_cookies_and_rejects_expired_access(monkeypatch):
    app = create_app(Settings(app_access_password=SecretStr("local-auth-fixture-only"),
                              app_origin="https://workspace.test", mock_latency_seconds=0))
    async with app.router.lifespan_context(app):
        async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="https://workspace.test") as client:
            assert (await client.get("/api/config")).json()["requires_login"]
            rejected = await client.post("/api/login", json={"password": "wrong-fixture"})
            assert rejected.status_code == 401 and "workspace_token" not in rejected.json()
            login = await client.post("/api/login", json={"password": "local-auth-fixture-only"})
            assert login.status_code == 200 and login.headers["cache-control"] == "no-store"
            assert "local-auth-fixture-only" not in login.text
            token = login.json()["workspace_token"]
            assert len(token) == 43
            assert not (await client.get("/api/config")).json()["requires_login"]
            # Reproduce a cross-site Space iframe that cannot retain cookies.
            client.cookies.clear()
            assert (await client.get("/api/config")).json()["requires_login"]
            assert (await client.post("/api/sessions", json=start_body())).status_code == 401
            assert (await client.post("/api/sessions", json=start_body(), headers={"X-Workspace-Token": "invalid"})).status_code == 401
            headers = {"X-Workspace-Token": token, "Origin": "https://workspace.test"}
            config = await client.get("/api/config", headers=headers)
            assert not config.json()["requires_login"] and token not in config.text
            assert (await client.post("/api/sessions", json=start_body(), headers={**headers, "Origin": "https://hostile.example"})).status_code == 403
            result = await client.post("/api/sessions", json=start_body(), headers=headers)
            assert result.status_code == 200 and token not in result.text
            session = result.json()
            assert (await client.post(f'/api/sessions/{session["session_id"]}/stop', headers={"X-Session-Token": session["session_token"]})).status_code == 200
            expired_time = time.monotonic() + 8 * 3600 + 1
            monkeypatch.setattr("backend.main.time", SimpleNamespace(monotonic=lambda: expired_time))
            assert (await client.get("/api/config", headers=headers)).json()["requires_login"]
            assert (await client.post("/api/sessions", json=start_body(), headers=headers)).status_code == 401


@pytest.mark.parametrize("update", [
    {"audio_b64": "invalid***"}, {"sample_count": 15999}, {"sample_rate": 48000},
    {"source_start_ms": 300}, {"source_end_ms": float("nan")},
    {"frames": [{"source_time_ms": 9000, "jpeg_b64": "/9j/2Q=="}]},
])
def test_invalid_audio_and_timestamps_rejected_before_model(update):
    with pytest.raises(ValidationError):
        AVChunk.model_validate({**media_chunk(), **update})


def test_600_second_budget_is_derived_from_native_cadence():
    settings = Settings()
    assert settings.chunk_budget(600000) >= 602
    assert settings.chunk_budget(600001) > settings.chunk_budget(600000)
    assert settings.max_duration_ms >= 600000
    assert settings.session_wall_timeout_seconds > 600 + settings.completion_timeout_seconds


def test_invalid_runtime_configuration_never_prints_secret_input():
    secret = "runtime-value-that-must-not-appear-in-error"
    with pytest.raises(ValidationError) as caught:
        Settings(app_mode="real", hf_token=secret, minicpm_service_url="https://invalid.example")
    assert secret not in str(caught.value)


@pytest.mark.asyncio
async def test_short_audio_tail_with_retained_visual_context_is_not_false_gap():
    from backend.schemas import StartRequest
    from backend.session import Session
    session = Session(StartRequest.model_validate(start_body(1037)), Settings(mock_latency_seconds=0))
    await session.start()
    await session.accept(AVChunk.model_validate(media_chunk()))
    tail = media_chunk(1, 592)
    tail["frames"] = []
    await session.accept(AVChunk.model_validate(tail))
    await session.complete()
    await session.completion_task
    assert session.coverage["complete"]
    assert session.coverage["processed_audio_ms"] == 1037
    assert "final_visual_hold" in session.diagnostics


@pytest.mark.asyncio
async def test_start_disconnect_cancels_native_preparation(monkeypatch):
    from backend.session import Session
    entered, cancelled, stopped = asyncio.Event(), asyncio.Event(), asyncio.Event()
    original_stop = Session.stop
    async def blocked_start(self):
        entered.set()
        try:
            await asyncio.sleep(30)
        except asyncio.CancelledError:
            cancelled.set()
            raise
    async def stop(self, reason="Manual Stop"):
        stopped.set()
        await original_stop(self, reason)
    monkeypatch.setattr(Session, "start", blocked_start)
    monkeypatch.setattr(Session, "stop", stop)
    app = create_app(Settings())
    import json
    body = json.dumps(start_body()).encode()
    sent = False
    async def receive():
        nonlocal sent
        if not sent:
            sent = True
            return {"type": "http.request", "body": body, "more_body": False}
        await entered.wait()
        return {"type": "http.disconnect"}
    messages = []
    async def send(message):
        messages.append(message)
    scope = {"type": "http", "asgi": {"version": "3.0"}, "http_version": "1.1", "method": "POST",
             "scheme": "http", "path": "/api/sessions", "raw_path": b"/api/sessions", "query_string": b"",
             "headers": [(b"content-type", b"application/json"), (b"content-length", str(len(body)).encode())],
             "client": ("127.0.0.1", 123), "server": ("127.0.0.1", 8000), "root_path": ""}
    async with app.router.lifespan_context(app):
        await asyncio.wait_for(app(scope, receive, send), 3)
        assert cancelled.is_set() and stopped.is_set()
        assert all(s.state == "stopped" for s in app.state.sessions.values())