test_server.py

"""Exercise the HTTP trust boundary without launching paid jobs or using ambient credentials."""

import sys
import time
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
from fastapi.testclient import TestClient

sys.path.insert(0, str(Path(__file__).parent))
import server  # noqa: E402


@pytest.fixture
def client():
    server.SESSIONS.clear()
    server.QUOTES.clear()
    with TestClient(server.app) as client:
        yield client


def login_client(client, username="alice"):
    response = client.get("/api/session")
    csrf = response.json()["csrf"]
    sid = client.cookies.get(server.COOKIE)
    server.SESSIONS[sid].update(user={"name": username}, token="never-return-this-token")
    return {"Origin": server.ORIGIN, "X-CSRF-Token": csrf}


def test_no_auth_token_is_ever_exposed_to_browser(client):
    login_client(client)
    response = client.get("/api/session")
    assert response.json()["user"]["name"] == "alice"
    assert "never-return-this-token" not in response.text
    assert "no-store" in response.headers["cache-control"]
    anonymous = TestClient(server.app).get("/api/session")
    assert "HttpOnly" in anonymous.headers["set-cookie"]


def test_anonymous_cannot_list_or_start_jobs(client):
    assert client.get("/api/jobs").status_code == 401
    response = client.post(
        "/api/jobs",
        headers={"Origin": server.ORIGIN},
        json={"quote_id": "a" * 24, "output": "alice/out", "budget": 5, "confirmed": True},
    )
    assert response.status_code == 401


@pytest.mark.parametrize("origin,csrf", [("https://evil.example", True), (server.ORIGIN, False)])
def test_cross_origin_and_missing_csrf_cannot_spend(client, monkeypatch, origin, csrf):
    headers = login_client(client)
    headers["Origin"] = origin
    if not csrf:
        headers.pop("X-CSRF-Token")
    submit = Mock()
    monkeypatch.setattr(server, "launch", submit)
    response = client.post(
        "/api/jobs",
        headers=headers,
        json={"quote_id": "a" * 24, "output": "alice/out", "budget": 5, "confirmed": True},
    )
    assert response.status_code == 403
    submit.assert_not_called()


def test_quote_is_user_bound_and_submission_is_idempotent(client, monkeypatch):
    headers = login_client(client)
    quote_id = "a" * 24
    server.QUOTES[quote_id] = {"created": time.time(), "quote_owner": "bob", "costs": {"hours": 1}}
    submit = Mock(return_value={"id": "job1", "owner": "alice", "output": "alice/out"})
    monkeypatch.setattr(server, "launch", submit)
    payload = {"quote_id": quote_id, "output": "alice/out", "budget": 5, "confirmed": True}
    assert client.post("/api/jobs", headers=headers, json=payload).status_code == 400
    submit.assert_not_called()
    server.QUOTES[quote_id]["quote_owner"] = "alice"
    assert client.post("/api/jobs", headers=headers, json=payload).json()["id"] == "job1"
    assert client.post("/api/jobs", headers=headers, json=payload).json()["id"] == "job1"
    submit.assert_called_once()


@pytest.mark.parametrize("output_label", ["alice/out", "out"])
def test_jobs_are_scoped_to_authenticated_person(client, monkeypatch, output_label):
    login_client(client)
    api = Mock()
    api.list_jobs.return_value = [
        SimpleNamespace(
            id="job1",
            status=SimpleNamespace(stage="RUNNING"),
            created_at="2026-09-25",
            labels={"annotation_output": output_label},
        )
    ]
    monkeypatch.setattr(server, "HfApi", lambda **kwargs: api)
    response = client.get("/api/jobs")
    assert response.status_code == 200
    api.list_jobs.assert_called_once_with(namespace="alice", labels={"app": "lerobot-annotation-studio"})
    assert "never-return-this-token" not in response.text
    assert response.json()[0]["output"] == "alice/out"


def test_oauth_callback_rejects_unmatched_state(client, monkeypatch):
    exchange = Mock()
    monkeypatch.setattr(server.requests, "post", exchange)
    response = client.get("/auth/callback?code=test&state=untrusted", follow_redirects=False)
    assert response.status_code == 307
    assert response.headers["location"] == "/?error=sign-in"
    exchange.assert_not_called()


def test_cancel_cannot_target_another_apps_job(client, monkeypatch):
    headers = login_client(client)
    api = Mock()
    api.inspect_job.return_value = SimpleNamespace(labels={"app": "another-app"})
    monkeypatch.setattr(server, "HfApi", lambda **kwargs: api)
    response = client.post("/api/jobs/job1/cancel", headers=headers, json={})
    assert response.status_code == 400
    api.inspect_job.assert_called_once_with(job_id="job1", namespace="alice")
    api.cancel_job.assert_not_called()


def test_prompt_catalog_and_validation_are_available_without_login(client):
    catalog = client.get("/api/prompts").json()
    assert len(catalog) == 7
    assert len({item["feature"] for item in catalog}) == 6
    item = catalog[0]
    edit = item["default"] + "\nKeep labels short."
    response = client.post(
        "/api/prompts/validate", headers={"Origin": server.ORIGIN}, json={"prompts": {item["key"]: edit}}
    )
    assert response.status_code == 200
    assert response.json() == {item["key"]: edit}
    invalid = client.post(
        "/api/prompts/validate",
        headers={"Origin": server.ORIGIN},
        json={"prompts": {item["key"]: "missing placeholders"}},
    )
    assert invalid.status_code == 400
    assert "placeholders" in invalid.json()["detail"]
    reset = client.post(
        "/api/prompts/validate",
        headers={"Origin": server.ORIGIN},
        json={"prompts": {item["key"]: item["default"]}},
    )
    assert reset.json() == {}


def test_model_catalog_and_unknown_model_rejection(client):
    response = client.get("/api/models")
    assert response.status_code == 200
    data = response.json()
    assert len(data["models"]) == 9
    assert data["default"] == "qwen38-27b"
    assert len([m for m in data["models"] if m["featured"]]) == 3
    assert data["unavailable"] == []
    assert all("deepseek" not in m["key"] for m in data["models"])
    response = client.post(
        "/api/estimate",
        headers={"Origin": server.ORIGIN},
        json={"dataset": "owner/data", "features": ["Subtasks"], "model": "unlisted-model"},
    )
    assert response.status_code == 400
    assert "supported vision model" in response.json()["detail"]


def test_camera_api_is_metadata_only_and_uses_the_current_user(client, monkeypatch):
    expected = {
        "dataset": "owner/robot",
        "revision": "a" * 40,
        "cameras": ["observation.images.top", "observation.images.wrist"],
    }
    discover = Mock(return_value=expected)
    monkeypatch.setattr(server, "read_cameras", discover)
    response = client.post("/api/cameras", headers={"Origin": server.ORIGIN}, json={"dataset": "owner/robot"})
    assert response.status_code == 200
    assert response.json() == expected
    discover.assert_called_once_with("owner/robot", False)
    assert not server.QUOTES
    headers = login_client(client)
    response = client.post("/api/cameras", headers=headers, json={"dataset": "owner/robot"})
    assert response.status_code == 200
    discover.assert_called_with("owner/robot", "never-return-this-token")
    assert "never-return-this-token" not in response.text


@pytest.mark.parametrize("visibility", [None, "private", "public"])
def test_launch_api_preserves_visibility_and_defaults_to_private(client, monkeypatch, visibility):
    headers = login_client(client)
    quote_id = "v" * 24
    server.QUOTES[quote_id] = {"created": time.time(), "quote_owner": "alice", "costs": {"hours": 1}}
    submit = Mock(return_value={"id": "job1", "owner": "alice", "output": "alice/out"})
    monkeypatch.setattr(server, "launch", submit)
    payload = {"quote_id": quote_id, "output": "alice/out", "budget": 5, "confirmed": True}
    if visibility is not None:
        payload["visibility"] = visibility
    assert client.post("/api/jobs", headers=headers, json=payload).status_code == 200
    assert submit.call_args.kwargs["visibility"] == (visibility or "private")


def test_launch_api_rejects_unknown_visibility(client, monkeypatch):
    headers = login_client(client)
    submit = Mock()
    monkeypatch.setattr(server, "launch", submit)
    response = client.post(
        "/api/jobs",
        headers=headers,
        json={
            "quote_id": "v" * 24,
            "output": "alice/out",
            "budget": 5,
            "confirmed": True,
            "visibility": "unknown",
        },
    )
    assert response.status_code == 422
    submit.assert_not_called()


def test_only_parquet_output_is_accepted(client):
    from studio import COPY, EXPORT

    assert server.EstimateInput(dataset="alice/data", features=["Subtasks"]).mode == COPY
    response = client.post(
        "/api/estimate",
        headers=login_client(client),
        json={"dataset": "alice/data", "features": ["Subtasks"], "mode": EXPORT},
    )
    assert response.status_code == 422