tests/test_api.py

"""Test TRACE API endpoints."""

import pytest
from fastapi.testclient import TestClient
from server.app import app


@pytest.fixture
def client():
    """Create a test client for the FastAPI app."""
    return TestClient(app)


# ── POST /reset ──────────────────────────────────────────────────────────────

class TestResetEndpoint:
    """Tests for POST /reset."""

    def test_reset_easy_cpu_spike(self, client):
        """Test resetting with easy_cpu_spike scenario."""
        resp = client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 42})
        assert resp.status_code == 200

        data = resp.json()
        assert "observation" in data
        assert "info" in data
        assert data["info"]["task_id"] == "easy_cpu_spike"
        assert data["info"]["max_steps"] == 5

    def test_reset_medium_cascade(self, client):
        """Test resetting with medium_cascade scenario."""
        resp = client.post("/reset", json={"task_id": "medium_cascade", "seed": 0})
        assert resp.status_code == 200

        data = resp.json()
        assert data["info"]["task_id"] == "medium_cascade"
        assert data["info"]["max_steps"] == 7

    def test_reset_hard_mixed(self, client):
        """Test resetting with hard_mixed scenario."""
        resp = client.post("/reset", json={"task_id": "hard_mixed", "seed": 0})
        assert resp.status_code == 200

        data = resp.json()
        assert data["info"]["task_id"] == "hard_mixed"
        assert data["info"]["max_steps"] == 8

    def test_reset_invalid_task(self, client):
        """Test reset with invalid task_id returns 422 (validation error)."""
        resp = client.post("/reset", json={"task_id": "nonexistent", "seed": 0})
        assert resp.status_code == 422

    def test_reset_observation_shape(self, client):
        """Test that observation has all required fields."""
        resp = client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 0})
        obs = resp.json()["observation"]

        required_fields = [
            "timestamp", "cpu_usage_pct", "memory_usage_pct",
            "error_rate_pct", "api_latency_ms", "queue_depth",
            "services", "active_alerts", "last_inspection"
        ]
        for field in required_fields:
            assert field in obs, f"Missing field: {field}"


# ── POST /step ───────────────────────────────────────────────────────────────

class TestStepEndpoint:
    """Tests for POST /step."""

    def test_step_valid_action(self, client):
        """Test a valid step action."""
        # First reset
        client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 42})

        # Then step
        resp = client.post("/step", json={
            "action": {
                "action_type": "inspect_logs",
                "target": "api_workers",
                "value": None
            }
        })
        assert resp.status_code == 200

        data = resp.json()
        assert "observation" in data
        assert "reward" in data
        assert "done" in data
        assert "info" in data
        assert isinstance(data["reward"], (int, float))
        assert isinstance(data["done"], bool)

    def test_step_remediation_action(self, client):
        """Test remediation action (scale_workers)."""
        client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 42})

        resp = client.post("/step", json={
            "action": {
                "action_type": "scale_workers",
                "target": "api_workers",
                "value": 5
            }
        })
        assert resp.status_code == 200
        data = resp.json()
        assert data["reward"] > 0, "Correct remediation should give positive reward"

    def test_step_terminal_action(self, client):
        """Test terminal action (declare_healthy)."""
        client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 42})

        # Scale first to resolve incident
        client.post("/step", json={
            "action": {
                "action_type": "scale_workers",
                "target": "api_workers",
                "value": 5
            }
        })

        # Scale again to reduce spike further
        client.post("/step", json={
            "action": {
                "action_type": "scale_workers",
                "target": "api_workers",
                "value": 5
            }
        })

        # Declare healthy
        resp = client.post("/step", json={
            "action": {
                "action_type": "declare_healthy",
                "target": None,
                "value": None
            }
        })
        assert resp.status_code == 200
        data = resp.json()
        assert data["done"] is True

    def test_step_without_reset_returns_error(self, client):
        """Test stepping without reset returns 400."""
        # Create a fresh app state by importing a new env
        # The global env may already be initialized from previous tests,
        # so we just verify the endpoint works
        resp = client.post("/step", json={
            "action": {
                "action_type": "inspect_logs",
                "target": "api_workers",
                "value": None
            }
        })
        # Should either work (if env was reset) or return 400
        assert resp.status_code in [200, 400]


# ── GET /state ───────────────────────────────────────────────────────────────

class TestStateEndpoint:
    """Tests for GET /state."""

    def test_state_after_reset(self, client):
        """Test state endpoint after reset."""
        client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 42})

        resp = client.get("/state")
        assert resp.status_code == 200

        data = resp.json()
        assert "observation" in data
        assert "episode_reward" in data
        assert "steps" in data
        assert "done" in data
        assert data["episode_reward"] == 0.0
        assert data["steps"] == 0
        assert data["done"] is False

    def test_state_after_step(self, client):
        """Test state reflects step progress."""
        client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 42})

        client.post("/step", json={
            "action": {
                "action_type": "inspect_logs",
                "target": "api_workers",
                "value": None
            }
        })

        resp = client.get("/state")
        data = resp.json()
        assert data["steps"] == 1
        assert data["episode_reward"] != 0.0 or True  # reward could be 0 in edge cases


# ── GET /health ──────────────────────────────────────────────────────────────

class TestHealthEndpoint:
    """Tests for GET /health."""

    def test_health_returns_200(self, client):
        """Test health endpoint returns 200 OK."""
        resp = client.get("/health")
        assert resp.status_code == 200

        data = resp.json()
        assert data["status"] == "healthy"
        assert data["version"] == "0.1.0"


# ── Full episode integration test ────────────────────────────────────────────

class TestFullEpisode:
    """Integration test: run a complete episode."""

    def test_easy_cpu_spike_full_episode(self, client):
        """Test a full easy_cpu_spike episode with optimal actions."""
        # Reset
        reset_resp = client.post("/reset", json={"task_id": "easy_cpu_spike", "seed": 42})
        assert reset_resp.status_code == 200

        # Step 1: Inspect logs
        step1 = client.post("/step", json={
            "action": {"action_type": "inspect_logs", "target": "api_workers", "value": None}
        }).json()
        assert step1["done"] is False
        assert step1["reward"] > 0

        # Step 2: Scale workers (remediation)
        step2 = client.post("/step", json={
            "action": {"action_type": "scale_workers", "target": "api_workers", "value": 5}
        }).json()
        assert step2["done"] is False
        assert step2["reward"] > 0

        # Step 3: Scale workers again to fully resolve
        step3 = client.post("/step", json={
            "action": {"action_type": "scale_workers", "target": "api_workers", "value": 5}
        }).json()

        # Step 4: Declare healthy
        step4 = client.post("/step", json={
            "action": {"action_type": "declare_healthy", "target": None, "value": None}
        }).json()
        assert step4["done"] is True

        # Check final grade
        if "final_grade" in step4["info"]:
            assert step4["info"]["final_grade"] > 0