tests/test_env.py

"""Phase 0 env verification suite — ENV-01..03, ENV-07, ENV-08.

ROADMAP Success Criteria literals:
  - `pytest tests/test_env.py::test_gold_answer_never_in_observation` passes (ENV-08)
  - All 4 termination reasons appear as named tests covering ENV-07 (a/b/c/d)
  - End-to-end HTTP roundtrip via FathomEnvClient proves ENV-01..03
"""
from __future__ import annotations

import time

import pytest
from fastapi.testclient import TestClient

from env.client import FathomEnvClient
from env.models import FathomAction, FathomObservation, FathomStepResult, TerminationReason


# ─── ENV-01: scaffold + HTTP roundtrip ────────────────────────────────────────

def test_healthz_returns_ok(fresh_test_client: TestClient) -> None:
    r = fresh_test_client.get("/healthz")
    assert r.status_code == 200
    assert r.json() == {"status": "ok"}


def test_reset_returns_observation_no_gold_field(fresh_test_client: TestClient) -> None:
    r = fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy"})
    assert r.status_code == 200, r.text
    obs_json = r.json()
    assert "gold_answer" not in obs_json, f"ENV-08 violation: gold_answer key in observation JSON: {obs_json}"
    assert "task_type" not in obs_json
    obs = FathomObservation.model_validate(obs_json)
    assert obs.tokens_remaining == 100_000
    assert obs.depth_max == 2
    assert obs.turns_left == 20


def test_step_roundtrip_print_1_plus_1(fresh_test_client: TestClient) -> None:
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy"})
    r = fresh_test_client.post("/step", json={"tool_name": "repl", "code": "print(1 + 1)"})
    assert r.status_code == 200
    body = r.json()
    assert "gold_answer" not in body.get("observation", {})
    sr = FathomStepResult.model_validate(body)
    assert "2" in sr.observation.stdout
    assert sr.done is False


# ─── ENV-08: gold-answer sealing (ROADMAP Success Criterion #3 literal) ───────

def test_gold_answer_never_in_observation(fresh_test_client: TestClient, fixture_rows: list[dict]) -> None:
    """For every fixture row, after reset() + several steps, no observation dict contains
    gold_answer or task_type as a key.

    This is the LITERAL test that ROADMAP Phase 0 Exit Gate Success Criterion #3 names.
    Key-presence is what's tested — the gold VALUE appearing inside context_preview is
    NOT a violation (the context legitimately contains it).
    """
    difficulties = ["easy", "medium", "hard"]
    for diff in difficulties:
        r = fresh_test_client.post("/reset", json={"seed": 0, "difficulty": diff})
        assert r.status_code == 200
        obs = r.json()
        assert "gold_answer" not in obs, f"diff={diff}: gold_answer leaked on reset: keys={list(obs.keys())}"
        assert "task_type" not in obs, f"diff={diff}: task_type leaked on reset: keys={list(obs.keys())}"

        for code in ("x = 1", "y = x + 2", "z = y * 3"):
            r2 = fresh_test_client.post("/step", json={"tool_name": "repl", "code": code})
            assert r2.status_code == 200
            body = r2.json()
            assert "gold_answer" not in body.get("observation", {}), f"step leak on diff={diff}, code={code}"
            assert "task_type" not in body.get("observation", {})


def test_state_endpoint_sanitized(fresh_test_client: TestClient) -> None:
    """ENV-08: /state returns a whitelist dict — no gold_answer, task_type, or raw context."""
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy"})
    r = fresh_test_client.get("/state")
    assert r.status_code == 200
    s = r.json()
    assert "gold_answer" not in s
    assert "task_type" not in s
    assert "context" not in s
    assert "episode_id" in s
    assert "step_count" in s
    assert "tokens_used_total" in s


# ─── ENV-07: 4 termination reasons ────────────────────────────────────────────

def test_termination_answer(fresh_test_client: TestClient) -> None:
    """ENV-07a — <answer>...</answer> emits done=True with reason=answer."""
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy"})
    r = fresh_test_client.post("/step", json={"tool_name": "repl", "code": "<answer>42</answer>"})
    sr = FathomStepResult.model_validate(r.json())
    assert sr.done is True
    assert sr.info.get("termination_reason") == TerminationReason.ANSWER.value
    assert sr.observation.answer_emitted is True
    assert sr.observation.return_val == "42"


def test_termination_max_steps(fresh_test_client: TestClient) -> None:
    """ENV-07b — exceeding max_steps terminates with reason=max_steps."""
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy", "max_steps": 2})
    r1 = fresh_test_client.post("/step", json={"tool_name": "repl", "code": "a = 1"})
    sr1 = FathomStepResult.model_validate(r1.json())
    assert sr1.done is False
    r2 = fresh_test_client.post("/step", json={"tool_name": "repl", "code": "b = 2"})
    sr2 = FathomStepResult.model_validate(r2.json())
    assert sr2.done is True
    assert sr2.info.get("termination_reason") == TerminationReason.MAX_STEPS.value


def test_termination_max_tokens(fresh_test_client: TestClient) -> None:
    """ENV-07c — exceeding max_tokens terminates with reason=max_tokens."""
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy", "max_tokens": 2})
    r = fresh_test_client.post(
        "/step",
        json={"tool_name": "repl", "code": "print('this is a long string designed to exceed two tokens easily')"},
    )
    sr = FathomStepResult.model_validate(r.json())
    assert sr.done is True, f"Expected done=True on max_tokens, got {sr}"
    assert sr.info.get("termination_reason") == TerminationReason.MAX_TOKENS.value


def test_termination_walltime(fresh_test_client: TestClient) -> None:
    """ENV-07d — walltime budget exceeded terminates with reason=walltime."""
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy", "walltime_budget_s": 0.3})
    time.sleep(0.5)
    r = fresh_test_client.post("/step", json={"tool_name": "repl", "code": "x = 1"})
    sr = FathomStepResult.model_validate(r.json())
    assert sr.done is True, f"Expected walltime termination, got {sr}"
    assert sr.info.get("termination_reason") == TerminationReason.WALLTIME.value


# ─── ENV-06: malformed code → structured error, done=False ────────────────────

def test_malformed_code_returns_error_observation_done_false(fresh_test_client: TestClient) -> None:
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy"})
    r = fresh_test_client.post("/step", json={"tool_name": "repl", "code": "1 / 0"})
    sr = FathomStepResult.model_validate(r.json())
    assert sr.done is False, "ENV-06: division by zero should NOT terminate the episode"
    combined = sr.observation.stderr + (sr.observation.return_val or "")
    assert "ZeroDivisionError" in combined


def test_malformed_syntax_returns_error_observation(fresh_test_client: TestClient) -> None:
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "easy"})
    r = fresh_test_client.post("/step", json={"tool_name": "repl", "code": "def (:"})
    sr = FathomStepResult.model_validate(r.json())
    assert sr.done is False
    assert sr.observation.stderr, "syntax error should surface in stderr"


# ─── ENV-03 / D-10: difficulty mapping ────────────────────────────────────────

def test_difficulty_mapping_medium_loads_multi_needle_row(fresh_test_client: TestClient) -> None:
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "medium"})
    r = fresh_test_client.get("/state")
    assert r.status_code == 200
    s = r.json()
    assert s["task_id"] == "fixture-0001-multi-needle", f"D-10 violation: got {s['task_id']}"


def test_difficulty_mapping_hard_loads_counting_row(fresh_test_client: TestClient) -> None:
    fresh_test_client.post("/reset", json={"seed": 0, "difficulty": "hard"})
    r = fresh_test_client.get("/state")
    s = r.json()
    assert s["task_id"] == "fixture-0002-counting"


def test_reset_by_task_id(fresh_test_client: TestClient) -> None:
    fresh_test_client.post("/reset", json={"seed": 0, "task_id": "fixture-0002-counting"})
    r = fresh_test_client.get("/state")
    s = r.json()
    assert s["task_id"] == "fixture-0002-counting"


# ─── ENV-01: live HTTP roundtrip via FathomEnvClient ──────────────────────────

def test_live_uvicorn_roundtrip(live_uvicorn_server: str) -> None:
    """End-to-end HTTP roundtrip matching ROADMAP Success Criterion #1.

    Spawns uvicorn in a subprocess, uses FathomEnvClient over the real HTTP wire,
    verifies reset() + step(print(1+1)) + state() all succeed with no gold leak.
    """
    with FathomEnvClient(base_url=live_uvicorn_server) as client:
        assert client.healthz() == {"status": "ok"}
        obs = client.reset(seed=0, difficulty="easy")
        assert isinstance(obs, FathomObservation)
        sr = client.step(FathomAction(tool_name="repl", code="print(1+1)"))
        assert isinstance(sr, FathomStepResult)
        assert "2" in sr.observation.stdout
        assert sr.done is False
        s = client.state()
        assert "gold_answer" not in s
        assert "task_type" not in s
        assert s.get("step_count") == 1