test_server.py

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

import json
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.mark.parametrize("written", [[], [0], [0, 1]])
def test_progress_separates_labels_from_viewable_exports(client, monkeypatch, written):
    login_client(client)
    job = SimpleNamespace(status=SimpleNamespace(stage="RUNNING"))
    report = {
        "status": "running",
        "expected_episodes": 2,
        "episodes": [{"episode_index": i, "status": "completed"} for i in range(2)],
        "written_episode_ids": written,
        "manifest": {"mode": server.COPY, "output": "alice/out"},
    }
    monkeypatch.setattr(server, "studio_output", lambda *args: (job, "alice/out"))
    monkeypatch.setattr(server, "read_output_file", lambda *args: json.dumps(report).encode())
    result = client.get("/api/jobs/job1/progress").json()
    assert result["completed"] == 2 and result["pending"] == 0
    assert result["exported"] == len(written)
    assert result["export_pending"] == 2 - len(written)
    assert bool(result["annotation_url"]) == bool(written)


@pytest.mark.parametrize("written", [[], [0], [2, 0]])
def test_hand_viewer_links_only_materialized_output_episodes(client, monkeypatch, written):
    login_client(client)
    monkeypatch.setenv("HAND_VISUALIZER_URL", "https://obsolete-viewer.example/")
    report = {
        "status": "running",
        "expected_episodes": 3,
        "episodes": [{"episode_index": 42, "status": "completed"}, {"episode_index": 9, "status": "failed"}],
        "written_episode_ids": written,
        "source_episode_map": [{"episode_index": 0, "source_episode_index": 42}],
        "manifest": {"mode": server.COPY, "output": "alice/hands", "hand_actions": {"enabled": True}},
    }
    job = SimpleNamespace(status=SimpleNamespace(stage="RUNNING"))
    monkeypatch.setattr(server, "studio_output", lambda *args: (job, "alice/hands"))
    monkeypatch.setattr(server, "read_output_file", lambda *args: json.dumps(report).encode())
    result = client.get("/api/jobs/job1/progress").json()
    assert result["hand_visualizer_url"] == (
        f"https://lerobot-visualize-dataset.hf.space/alice/hands/{min(written)}?tab=hands"
        if written
        else None
    )
    report["manifest"].pop("hand_actions")
    assert client.get("/api/jobs/job1/progress").json()["hand_visualizer_url"] is None


@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_require_login(client):
    headers = login_client(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=headers, json={"prompts": {item["key"]: edit}})
    assert response.status_code == 200
    assert response.json() == {item["key"]: edit}
    invalid = client.post(
        "/api/prompts/validate",
        headers=headers,
        json={"prompts": {item["key"]: "missing placeholders"}},
    )
    assert invalid.status_code == 400
    assert "placeholders" in invalid.json()["detail"]
    reset = client.post(
        "/api/prompts/validate",
        headers=headers,
        json={"prompts": {item["key"]: item["default"]}},
    )
    assert reset.json() == {}


def test_model_catalog_and_unknown_model_rejection(client):
    headers = login_client(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=headers,
        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 == 401
    discover.assert_not_called()
    assert not server.QUOTES
    headers = login_client(client)
    response = client.post("/api/cameras", headers=headers, json={"dataset": "owner/robot"})
    assert response.status_code == 200
    assert response.json() == expected
    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


@pytest.mark.parametrize(
    "method,path",
    [
        ("GET", "/api/models"),
        ("GET", "/api/prompts"),
        ("GET", "/api/hand-pipeline"),
        ("GET", "/api/slam-pipeline"),
        ("GET", "/api/jobs"),
        ("GET", "/api/jobs/other/progress"),
        ("GET", "/api/jobs/other/human-videos"),
        ("GET", "/api/jobs/other/media?path=secret"),
        ("POST", "/api/cameras"),
        ("POST", "/api/estimate"),
        ("POST", "/api/prompts/validate"),
        ("POST", "/api/jobs"),
        ("POST", "/api/jobs/other/cancel"),
        ("POST", "/api/future-route"),
    ],
)
def test_all_functional_apis_fail_closed_without_login(client, monkeypatch, method, path):
    monkeypatch.setenv("HF_TOKEN", "owner-credential-must-never-be-used")
    hub = Mock(side_effect=AssertionError("Anonymous request reached HF"))
    monkeypatch.setattr(server, "HfApi", hub)
    monkeypatch.setattr(server, "read_cameras", hub)
    monkeypatch.setattr(server, "prepare", hub)
    monkeypatch.setattr(server, "launch", hub)
    response = client.request(method, path, headers={"Origin": server.ORIGIN}, json={})
    assert response.status_code == 401
    assert response.headers["cache-control"] == "no-store"
    assert "owner-credential" not in response.text
    hub.assert_not_called()


@pytest.mark.parametrize("path", ["/api/cameras", "/api/estimate", "/api/prompts/validate"])
def test_non_launch_posts_also_require_csrf(client, path):
    headers = login_client(client)
    headers.pop("X-CSRF-Token")
    assert client.post(path, headers=headers, json={}).status_code == 403


@pytest.mark.parametrize("owner", [None, "bob"])
def test_anonymous_or_other_users_quotes_cannot_be_adopted(client, monkeypatch, owner):
    headers = login_client(client)
    server.QUOTES["a" * 24] = {"created": time.time(), "quote_owner": owner}
    submit = Mock()
    monkeypatch.setattr(server, "launch", submit)
    assert (
        client.post(
            "/api/jobs",
            headers=headers,
            json={
                "quote_id": "a" * 24,
                "output": "alice/out",
                "budget": 5,
                "confirmed": True,
            },
        ).status_code
        == 400
    )
    submit.assert_not_called()


def test_expired_session_cannot_use_ambient_credentials(client, monkeypatch):
    headers = login_client(client)
    server.SESSIONS[client.cookies.get(server.COOKIE)]["expires"] = time.time() - 1
    submit = Mock()
    monkeypatch.setattr(server, "launch", submit)
    monkeypatch.setenv("HF_TOKEN", "space-owner-token")
    assert client.post("/api/jobs", headers=headers, json={}).status_code == 401
    submit.assert_not_called()


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


def test_media_cannot_read_outputs_in_another_namespace(client, monkeypatch):
    login_client(client)
    api = Mock()
    api.inspect_job.return_value = SimpleNamespace(
        labels={
            "app": "lerobot-annotation-studio",
            "annotation_output": "bob/private",
        }
    )
    factory = Mock(return_value=api)
    read = Mock()
    monkeypatch.setattr(server, "HfApi", factory)
    monkeypatch.setattr(server, "read_output_file", read)
    response = client.get(
        "/api/jobs/job1/media?path=annotation_studio/human_videos/episode_000000/segment_000.mp4"
    )
    assert response.status_code == 404
    factory.assert_called_once_with(token="never-return-this-token")
    api.inspect_job.assert_called_once_with(job_id="job1", namespace="alice")
    read.assert_not_called()


def test_unknown_api_is_not_a_cached_spa_document(client):
    login_client(client)
    assert client.get("/api/not-a-route").status_code == 404
    response = client.get("/")
    assert response.headers["cache-control"] == "no-store"
    assert "script-src 'self'" in response.headers["content-security-policy"]
    assert "object-src 'none'" in response.headers["content-security-policy"]
    assert "Open Annotation Studio" in response.text


def test_oauth_callback_rotates_session_and_cannot_replay(client, monkeypatch):
    monkeypatch.setenv("OAUTH_CLIENT_ID", "app-id")
    monkeypatch.setenv("OAUTH_CLIENT_SECRET", "app-secret")
    monkeypatch.setattr(server, "HOST", "studio.hf.space")
    # The live origin is HTTPS. Do not force Secure cookies into an HTTP test client.
    with TestClient(server.app, base_url="https://studio.hf.space") as secure_client:
        response = secure_client.get("/auth/login", follow_redirects=False)
        old_sid = secure_client.cookies.get(server.COOKIE)
        state = server.SESSIONS[old_sid]["oauth_state"]
        assert "Secure" in response.headers["set-cookie"]
        exchange = Mock(
            return_value=SimpleNamespace(
                raise_for_status=lambda: None,
                json=lambda: {"access_token": "oauth-user-token", "expires_in": 86400},
            )
        )
        api = Mock()
        api.whoami.return_value = {"name": "alice"}
        factory = Mock(return_value=api)
        monkeypatch.setattr(server.requests, "post", exchange)
        monkeypatch.setattr(server, "HfApi", factory)
        response = secure_client.get(f"/auth/callback?code=one-use&state={state}", follow_redirects=False)
        assert response.headers["location"] == "/"
        assert old_sid not in server.SESSIONS
        assert secure_client.cookies.get(server.COOKIE) != old_sid
        factory.assert_called_once_with(token="oauth-user-token")
        assert "oauth-user-token" not in response.text
        response = secure_client.get(f"/auth/callback?code=one-use&state={state}", follow_redirects=False)
        assert response.headers["location"] == "/?error=sign-in"
        exchange.assert_called_once()


def test_raw_grabette_discovery_uses_session_token(client, monkeypatch):
    headers = login_client(client)
    reader = Mock(return_value={"cameras": ["right"], "source": "raw-grabette"})
    monkeypatch.setattr(server, "read_raw_cameras", reader)
    response = client.post(
        "/api/cameras", headers=headers, json={"dataset": "alice/raw", "raw_grabette": True}
    )
    assert response.status_code == 200
    reader.assert_called_once_with("alice/raw", "never-return-this-token")


def test_removed_generic_slam_parameters_rejected(client):
    headers = login_client(client)
    response = client.post(
        "/api/estimate",
        headers=headers,
        json={
            "dataset": "alice/raw",
            "features": ["Grabette SLAM"],
            "slam": {"calibration_mode": "manual", "intrinsics": {"fx": 400}},
        },
    )
    assert response.status_code == 422