tests/test_api.py

import json

import pytest
from conftest import GOOD_SCENE
from fastapi.testclient import TestClient

from api import keys, review
from api.auth import Allowlist
from api.llm import LLMError, ScriptedLLM
from api.logstore import LogStore
from api.main import Services, create_app


class FailingLLM:
    model = "z-ai/glm-5.2"

    def complete(self, messages, temperature=0.2, model=None):
        raise LLMError("LLM request failed: HTTP 503")


@pytest.fixture
def env(settings, store):
    key = keys.add(settings.testers_file, "alice")
    llm = ScriptedLLM([])
    services = Services(
        settings=settings,
        store=store,
        logs=LogStore(settings.log_db),
        allowlist=Allowlist(settings.testers_file),
        llm_factory=lambda: llm,
    )
    return TestClient(create_app(services)), key, llm, services


def test_health_is_public(env):
    client, *_ = env
    r = client.get("/health")
    assert r.status_code == 200 and r.json()["model"] == "z-ai/glm-5.2"
    assert r.json()["index"]["mujoco"] == 3


@pytest.mark.parametrize("headers", [{}, {"X-API-Key": "rd_wrong"}, {"Authorization": "Bearer nope"}])
def test_rejects_unknown_keys(env, headers):
    client, *_ = env
    r = client.post("/simulate/generate", json={"description": "a floor with a box"}, headers=headers)
    assert r.status_code == 401


def test_simulate_end_to_end_and_logged(env):
    client, key, llm, services = env
    llm.responses[:] = [f"```xml\n{GOOD_SCENE}\n```"]
    r = client.post(
        "/simulate/generate",
        json={"description": "G1 beside a 15 degree slope and a crate"},
        headers={"Authorization": f"Bearer {key}"},
    )
    assert r.status_code == 200, r.text
    body = r.json()
    assert body["valid"] and body["mjcf"] == GOOD_SCENE and body["attempts"] == 1
    assert "mujoco_menagerie/unitree_g1" in body["how_to_run"]

    row = services.logs.get(body["id"])
    assert row["tester"] == "alice" and row["kind"] == "simulate" and row["valid"] is True
    assert row["request"]["description"].startswith("G1 beside")
    assert row["attempts"][0]["raw_output"].startswith("```xml")
    assert row["sources"] == body["sources"]


def test_llm_failure_is_logged_and_reported(env):
    client, key, _, services = env
    services.llm_factory = lambda: FailingLLM()
    r = client.post("/simulate/generate", json={"description": "a floor with a box"}, headers={"X-API-Key": key})
    assert r.status_code == 502 and "log id" in r.json()["detail"]
    (row,) = services.logs.list()
    assert row["error"].startswith("LLM request failed")


def test_request_validation(env):
    client, key, *_ = env
    h = {"X-API-Key": key}
    assert client.post("/simulate/generate", json={"description": "short"}, headers=h).status_code == 422
    assert (
        client.post("/simulate/generate", json={"description": "x" * 20, "robot": "atlas"}, headers=h).status_code
        == 422
    )


def test_feedback_only_for_own_generations(env):
    client, key, llm, services = env
    llm.responses[:] = [f"```xml\n{GOOD_SCENE}\n```"]
    gen_id = client.post(
        "/simulate/generate", json={"description": "G1 and a crate on the floor"}, headers={"X-API-Key": key}
    ).json()["id"]
    r = client.post(
        f"/feedback/{gen_id}", json={"rating": "works", "note": "loads in viewer"}, headers={"X-API-Key": key}
    )
    assert r.status_code == 200
    assert services.logs.get(gen_id)["tester_rating"] == "works"

    bob = keys.add(services.settings.testers_file, "bob")
    r = client.post(f"/feedback/{gen_id}", json={"rating": "broken"}, headers={"X-API-Key": bob})
    assert r.status_code == 404


class TestKeys:
    def test_add_list_revoke(self, tmp_path, capsys):
        path = tmp_path / "testers.yaml"
        key = keys.add(path, "alice", "robotics lab")
        allow = Allowlist(path)
        assert allow.identify(key).name == "alice"
        assert key not in path.read_text()  # only the hash is stored
        with pytest.raises(ValueError, match="already exists"):
            keys.add(path, "alice")
        assert keys.revoke(path, "alice") == 1
        assert allow.identify(key) is None  # picked up without a restart
        assert keys.main(["--file", str(path), "list"]) == 0
        assert "revoked  alice" in capsys.readouterr().out

    def test_env_keys(self, tmp_path, monkeypatch):
        monkeypatch.setenv("ROSDIFF_API_KEYS", "dev:secret123, other:abc")
        allow = Allowlist(tmp_path / "missing.yaml")
        assert allow.identify("secret123").name == "dev"
        assert allow.identify("abc").name == "other"
        assert allow.identify("") is None


class TestReview:
    def test_list_show_mark_export(self, tmp_path, capsys):
        db = tmp_path / "log.sqlite3"
        logs = LogStore(db)
        sim = logs.record(
            tester="alice",
            kind="simulate",
            request={"description": "G1 on stairs"},
            output=GOOD_SCENE,
            valid=True,
            validation={"valid": True},
            attempts=[{"raw_output": "x"}],
        )
        code = logs.record(
            tester="bob",
            kind="code",
            request={"request": "lidar node"},
            output=json.dumps({"pkg/package.xml": "<package/>"}),
            valid=False,
        )
        run = lambda *a: review.main(["--db", str(db), *a])  # noqa: E731

        assert run("list") == 0
        out = capsys.readouterr().out
        assert sim in out and code in out and "G1 on stairs" in out

        assert run("mark", sim, "correct", "--note", "stairs look right") == 0
        assert run("list", "--unreviewed") == 0
        out = capsys.readouterr().out
        assert sim not in out.split("correct")[-1] and code in out

        assert run("show", code) == 0
        assert "=== pkg/package.xml" in capsys.readouterr().out

        assert run("files", sim, str(tmp_path / "out")) == 0
        assert (tmp_path / "out" / "scene_rosdiff.xml").read_text() == GOOD_SCENE

        assert run("summary") == 0
        assert "reviewed correct 1/1" in capsys.readouterr().out

        assert run("export", str(tmp_path / "all.jsonl")) == 0
        assert len((tmp_path / "all.jsonl").read_text().splitlines()) == 2

    def test_files_refuses_path_traversal(self, tmp_path, capsys):
        logs = LogStore(tmp_path / "log.sqlite3")
        gen = logs.record(tester="x", kind="code", request={}, output=json.dumps({"../../evil.py": "x"}))
        review.main(["--db", str(tmp_path / "log.sqlite3"), "files", gen, str(tmp_path / "out")])
        assert not (tmp_path / "evil.py").exists()
        assert "unsafe path" in capsys.readouterr().err


def test_code_generate_direct_backend(env):
    from ros_fixtures import INTERFACES, as_file_blocks, python_package

    from api.ros_check import load_interfaces

    client, key, llm, services = env
    services.settings.data_dir.mkdir(parents=True, exist_ok=True)
    services.settings.interfaces_file.write_text(json.dumps(INTERFACES))
    load_interfaces.cache_clear()
    llm.responses[:] = [as_file_blocks(python_package())]
    r = client.post(
        "/code/generate",
        headers={"X-API-Key": key},
        json={
            "request": "Stop the robot when LaserScan sees something closer than 0.5 m",
            "distro": "jazzy",
            "package_name": "safety_stop",
        },
    )
    assert r.status_code == 200, r.text
    body = r.json()
    assert body["valid"] and body["backend"] == "direct" and "safety_stop/package.xml" in body["files"]
    assert "source /opt/ros/jazzy/setup.bash" in body["how_to_build"]
    row = services.logs.get(body["id"])
    assert row["kind"] == "code" and row["valid"] is True
    assert json.loads(row["output"]) == body["files"]
    assert row["request"]["distro"] == "jazzy"


def test_code_generate_rejects_bad_input(env):
    client, key, *_ = env
    h = {"X-API-Key": key}
    assert (
        client.post("/code/generate", json={"request": "a talker node", "distro": "rolling"}, headers=h).status_code
        == 422
    )
    assert (
        client.post(
            "/code/generate", json={"request": "a talker node", "package_name": "Bad-Name"}, headers=h
        ).status_code
        == 422
    )


def test_coding_model_is_separate(settings):
    from dataclasses import replace

    from api.opencode_runner import opencode_config

    s = replace(settings, llm_model="z-ai/glm-5.2", code_llm_model="some/coder")
    cfg = opencode_config(s)
    models = list(cfg["provider"]["llm"]["models"])
    assert cfg["model"] == "llm/some/coder" and models[0] == "some/coder" and "deepseek/deepseek-v4-flash" in models
    assert cfg["provider"]["llm"]["options"]["baseURL"] == s.llm_base_url


def test_code_direct_backend_uses_code_model(env):
    from ros_fixtures import INTERFACES, as_file_blocks, python_package

    from api.llm import ScriptedLLM
    from api.ros_check import load_interfaces

    client, key, llm, services = env
    services.settings.data_dir.mkdir(parents=True, exist_ok=True)
    services.settings.interfaces_file.write_text(json.dumps(INTERFACES))
    load_interfaces.cache_clear()
    coder = ScriptedLLM([as_file_blocks(python_package())], model="some/coder")
    services.code_llm_factory = lambda: coder
    r = client.post(
        "/code/generate", headers={"X-API-Key": key}, json={"request": "Stop on LaserScan closer than 0.5 m"}
    )
    body = r.json()
    # the router picks the model (a small node -> the low-cost tier); the code client is the one used
    assert r.status_code == 200 and coder.calls and not llm.calls
    assert body["model"] == coder.models[0] == "deepseek/deepseek-v4-flash" and body["routing"]["tier"] == "low"
    health = client.get("/health").json()
    assert health["model"] == services.settings.llm_model and "code_model" in health