test_env.py

"""
tests/test_env.py — Validates the CustomerSupportTriage environment.

Run: pytest tests/ -v
"""
import sys
sys.path.insert(0, "server")

import pytest
from models import (
    BatchTriageAction, Department, Priority, ResetRequest, TriageAction,
)
from env import CustomerSupportTriageEnv


@pytest.fixture
def env():
    return CustomerSupportTriageEnv()


# ── reset ──────────────────────────────────────────────────────────────────────

def test_reset_easy(env):
    result = env.reset(ResetRequest(task="easy", seed=0))
    assert result.task == "easy"
    assert len(result.observation.queue) == 5
    assert result.observation.step_number == 0
    assert result.observation.processed == 0


def test_reset_medium(env):
    result = env.reset(ResetRequest(task="medium", seed=0))
    assert len(result.observation.queue) == 10


def test_reset_hard(env):
    result = env.reset(ResetRequest(task="hard", seed=0))
    assert len(result.observation.queue) == 15


def test_reset_clears_state(env):
    env.reset(ResetRequest(task="easy", seed=0))
    obs = env.reset(ResetRequest(task="medium", seed=0)).observation
    assert obs.processed == 0
    assert obs.task_name == "medium"


# ── step ───────────────────────────────────────────────────────────────────────

def _make_action(ticket_id: str) -> TriageAction:
    return TriageAction(
        ticket_id=ticket_id,
        priority=Priority.MEDIUM,
        department=Department.GENERAL,
        response="Thank you for contacting us. We are reviewing your request and will respond shortly.",
        needs_human=False,
        reasoning="test",
    )


def test_step_advances_queue(env):
    result = env.reset(ResetRequest(task="easy", seed=0))
    first_ticket_id = result.observation.queue[0].ticket_id

    step_result = env.step(BatchTriageAction(actions=[_make_action(first_ticket_id)]))
    assert step_result.observation.processed == 1
    assert len(step_result.observation.queue) == 4


def test_step_reward_in_range(env):
    result = env.reset(ResetRequest(task="easy", seed=0))
    tids = [t.ticket_id for t in result.observation.queue]
    actions = [_make_action(tid) for tid in tids]
    step_result = env.step(BatchTriageAction(actions=actions))
    assert 0.0 <= step_result.reward <= 1.0


def test_step_done_when_queue_empty(env):
    result = env.reset(ResetRequest(task="easy", seed=0))
    tids = [t.ticket_id for t in result.observation.queue]
    step_result = env.step(BatchTriageAction(actions=[_make_action(t) for t in tids]))
    assert step_result.done is True
    assert len(step_result.observation.queue) == 0


def test_step_raises_when_done(env):
    env.reset(ResetRequest(task="easy", seed=0))
    obs = env.reset(ResetRequest(task="easy", seed=0)).observation
    tids = [t.ticket_id for t in obs.queue]
    env.step(BatchTriageAction(actions=[_make_action(t) for t in tids]))
    with pytest.raises(RuntimeError):
        env.step(BatchTriageAction(actions=[_make_action(tids[0])]))


def test_partial_step(env):
    result = env.reset(ResetRequest(task="easy", seed=0))
    first = result.observation.queue[0].ticket_id
    step_result = env.step(BatchTriageAction(actions=[_make_action(first)]))
    assert step_result.done is False
    assert step_result.observation.processed == 1


# ── state ──────────────────────────────────────────────────────────────────────

def test_state_after_reset(env):
    env.reset(ResetRequest(task="easy", seed=0))
    state = env.state()
    assert state.task_name == "easy"
    assert state.total_tickets == 5
    assert state.pending == 5
    assert state.processed == 0
    assert state.done is False


def test_state_updates_after_step(env):
    result = env.reset(ResetRequest(task="easy", seed=0))
    tid = result.observation.queue[0].ticket_id
    env.step(BatchTriageAction(actions=[_make_action(tid)]))
    state = env.state()
    assert state.processed == 1
    assert state.pending == 4
    assert len(state.reward_history) == 1


# ── grader quality ─────────────────────────────────────────────────────────────

def test_perfect_action_scores_high():
    """Correct priority+department+response should score ≥ 0.7."""
    from tasks import TASKS, grade_action, build_ticket

    for raw in TASKS["easy"].tickets:
        ticket = build_ticket(raw)
        action = TriageAction(
            ticket_id  = ticket.ticket_id,
            priority   = ticket._true_priority,
            department = ticket._true_department,
            response   = ticket._true_response,
            needs_human= ticket._true_department == Department.ESCALATION
                         or ticket._true_priority == Priority.URGENT,
            reasoning  = "ground truth",
        )
        info = grade_action(action, ticket)
        assert info.total >= 0.7, f"Ticket {ticket.ticket_id}: score={info.total}"


def test_wrong_action_scores_low():
    """Wrong everything should score < 0.4."""
    from tasks import TASKS, grade_action, build_ticket

    raw    = TASKS["easy"].tickets[0]   # T001: urgent / technical
    ticket = build_ticket(raw)
    action = TriageAction(
        ticket_id  = ticket.ticket_id,
        priority   = Priority.LOW,
        department = Department.SHIPPING,
        response   = "ok",              # very short → penalised
        needs_human= False,
        reasoning  = "wrong",
    )
    info = grade_action(action, ticket)
    assert info.total < 0.4, f"Expected low score, got {info.total}"


def test_grader_deterministic():
    """Same input → same output."""
    from tasks import TASKS, grade_action, build_ticket

    raw    = TASKS["easy"].tickets[0]
    ticket = build_ticket(raw)
    action = TriageAction(
        ticket_id  = ticket.ticket_id,
        priority   = ticket._true_priority,
        department = ticket._true_department,
        response   = "We are looking into your issue and will respond shortly.",
        needs_human= True,
        reasoning  = "test",
    )
    score1 = grade_action(action, ticket).total
    score2 = grade_action(action, ticket).total
    assert score1 == score2


def test_all_tasks_have_tickets():
    from tasks import TASKS
    for name, task in TASKS.items():
        assert len(task.tickets) >= 3, f"{name} needs ≥3 tickets"


def test_reward_range_full_episode():
    """Run a full easy episode and ensure reward ∈ [0,1]."""
    env = CustomerSupportTriageEnv()
    result = env.reset(ResetRequest(task="easy", seed=7))
    tids   = [t.ticket_id for t in result.observation.queue]
    sr     = env.step(BatchTriageAction(actions=[_make_action(t) for t in tids]))
    assert 0.0 <= sr.reward <= 1.0
    assert sr.done is True