tests/test_api.py

from fastapi.testclient import TestClient

from server.app import app


client = TestClient(app)


def test_reset_endpoint_contract() -> None:
    response = client.post("/reset", json={"task": "easy_injection"})
    payload = response.json()

    assert response.status_code == 200
    assert "observation" in payload
    assert payload["observation"]["task_id"] == "easy_injection"
    assert payload["observation"]["step_id"] == 1


def test_reset_endpoint_without_body_uses_default_task() -> None:
    response = client.post("/reset")
    payload = response.json()

    assert response.status_code == 200
    assert payload["observation"]["task_id"] == "easy_injection"


def test_step_endpoint_contract() -> None:
    client.post("/reset", json={"task": "easy_injection"})
    response = client.post(
        "/step",
        json={
            "action_type": "BLOCK",
            "justification": "Prompt asks for secrets and must be denied.",
        },
    )
    payload = response.json()

    assert response.status_code == 200
    assert set(payload.keys()) == {"observation", "reward", "done", "info"}
    assert "expected_action" in payload["info"]
    assert "reward_breakdown" in payload["info"]


def test_state_endpoint_contract() -> None:
    client.post("/reset", json={"task": "hard_rugpull"})
    response = client.get("/state")
    payload = response.json()

    assert response.status_code == 200
    assert payload["task_name"] == "hard_rugpull"
    assert "available_tasks" in payload
    assert "metrics" in payload