tests/test_environment.py

from models import ActionType, B2BSupportPayload, B2BSupportTriageAction
from server.support_triage_environment import B2BSupportTriageEnvironment


def test_reset_initializes_clean_state() -> None:
    env = B2BSupportTriageEnvironment()
    obs = env.reset(task_id="easy", seed=11)

    assert obs.task_id == "easy"
    assert obs.step_index == 0
    assert 0.0 < obs.progress_score < 1.0
    assert env.state.step_count == 0
    assert env.state.applied_decisions == {}


def test_invalid_ticket_penalty() -> None:
    env = B2BSupportTriageEnvironment()
    env.reset(task_id="easy")

    obs = env.step(
        B2BSupportTriageAction(
            action_type=ActionType.CLASSIFY,
            ticket_id="WRONG",
            payload=B2BSupportPayload(category="billing"),
        )
    )

    assert obs.reward < 0
    assert obs.last_action_error is not None
    assert "ticket_id_mismatch" in obs.last_action_error


def test_hard_task_can_reach_full_score() -> None:
    env = B2BSupportTriageEnvironment()
    reset_obs = env.reset(task_id="hard")
    ticket_id = reset_obs.visible_ticket.ticket_id

    env.step(
        B2BSupportTriageAction(
            action_type=ActionType.CLASSIFY,
            ticket_id=ticket_id,
            payload=B2BSupportPayload(category="security"),
        )
    )
    env.step(
        B2BSupportTriageAction(
            action_type=ActionType.SET_PRIORITY,
            ticket_id=ticket_id,
            payload=B2BSupportPayload(priority="urgent"),
        )
    )
    env.step(
        B2BSupportTriageAction(
            action_type=ActionType.ROUTE,
            ticket_id=ticket_id,
            payload=B2BSupportPayload(
                route_queue="security-incident-response",
                sla_minutes=120,
                escalate=True,
            ),
        )
    )
    env.step(
        B2BSupportTriageAction(
            action_type=ActionType.DRAFT_REPLY,
            ticket_id=ticket_id,
            payload=B2BSupportPayload(
                reply_text=(
                    "We have escalated this to our security team. "
                    "This issue is escalated. Please reset your API key now. "
                    "We will provide an update in 2 hours."
                )
            ),
        )
    )
    final_obs = env.step(
        B2BSupportTriageAction(
            action_type=ActionType.SUBMIT,
            ticket_id=None,
            payload=B2BSupportPayload(),
        )
    )

    assert final_obs.done is True
    assert 0.99 <= final_obs.progress_score < 1.0


def test_episode_ends_at_max_steps() -> None:
    env = B2BSupportTriageEnvironment()
    obs = env.reset(task_id="easy")
    ticket_id = obs.visible_ticket.ticket_id

    final_obs = obs
    for _ in range(obs.max_steps):
        final_obs = env.step(
            B2BSupportTriageAction(
                action_type=ActionType.CLASSIFY,
                ticket_id=ticket_id,
                payload=B2BSupportPayload(category="billing"),
            )
        )

    assert final_obs.done is True
    assert final_obs.last_action_error == "max_steps_reached"