tests/test_environment.py

"""
Module: test_environment.py
Purpose: Test core environment behavior — reset, step, state.
Part of: Medical Triage Assistant — OpenEnv Round 1
Author: Team Squirrel
"""

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import pytest
from triage_flow.environment import TriageEnvironment
from models import TriageAction, ActionType, PriorityLevel, InfoField


class TestReset:
    """Test reset() behavior."""

    def test_reset_returns_observation(self):
        """reset() should return a TriageObservation."""
        env = TriageEnvironment()
        obs = env.reset(task_name="basic-triage")
        assert obs.done is False
        assert obs.current_patient is not None
        assert obs.queue_length > 0

    def test_reset_basic_triage_has_3_patients(self):
        """Basic triage should have 3 patients."""
        env = TriageEnvironment()
        obs = env.reset(task_name="basic-triage")
        assert obs.queue_length == 3

    def test_reset_clears_state(self):
        """reset() should clear all previous state."""
        env = TriageEnvironment()
        obs = env.reset(task_name="basic-triage")

        # Take an action
        action = TriageAction(
            action_type=ActionType.ASSIGN_PRIORITY,
            patient_id="P001",
            priority_level=PriorityLevel.IMMEDIATE
        )
        env.step(action)

        # Reset should clear
        obs = env.reset(task_name="basic-triage")
        assert env.state.step_count == 0
        assert env.state.assignments == {}

    def test_reset_all_tasks(self):
        """All 3 tasks should reset successfully."""
        env = TriageEnvironment()
        for task in ["basic-triage", "incomplete-records-triage", "mass-casualty-triage"]:
            obs = env.reset(task_name=task)
            assert obs.done is False
            assert obs.current_patient is not None


class TestStep:
    """Test step() behavior."""

    def test_step_valid_action(self):
        """step() with valid action should return updated observation."""
        env = TriageEnvironment()
        env.reset(task_name="basic-triage")

        action = TriageAction(
            action_type=ActionType.ASSIGN_PRIORITY,
            patient_id="P001",
            priority_level=PriorityLevel.IMMEDIATE,
        )
        obs = env.step(action)
        assert obs.reward is not None
        assert obs.step_number == 1

    def test_step_correct_assignment_positive_reward(self):
        """Correct priority assignment should give positive reward."""
        env = TriageEnvironment()
        env.reset(task_name="basic-triage")

        # P001 ground truth is IMMEDIATE
        action = TriageAction(
            action_type=ActionType.ASSIGN_PRIORITY,
            patient_id="P001",
            priority_level=PriorityLevel.IMMEDIATE,
        )
        obs = env.step(action)
        assert obs.reward > 0

    def test_step_incorrect_assignment_negative_reward(self):
        """Incorrect priority assignment should give negative or zero reward."""
        env = TriageEnvironment()
        env.reset(task_name="basic-triage")

        # P001 ground truth is IMMEDIATE, assigning NON_URGENT is wrong
        action = TriageAction(
            action_type=ActionType.ASSIGN_PRIORITY,
            patient_id="P001",
            priority_level=PriorityLevel.NON_URGENT,
        )
        obs = env.step(action)
        assert obs.reward < 0

    def test_step_counts_increment(self):
        """Step count should increment with each step."""
        env = TriageEnvironment()
        env.reset(task_name="basic-triage")

        action = TriageAction(action_type=ActionType.ADVANCE_QUEUE)
        obs = env.step(action)
        assert obs.step_number == 1

        obs = env.step(action)
        assert obs.step_number == 2

    def test_done_when_all_assigned(self):
        """Episode should end when all patients are assigned."""
        env = TriageEnvironment()
        env.reset(task_name="basic-triage")

        # Assign all 3 patients
        for pid, priority in [
            ("P001", PriorityLevel.IMMEDIATE),
            ("P002", PriorityLevel.NON_URGENT),
            ("P003", PriorityLevel.URGENT),
        ]:
            action = TriageAction(
                action_type=ActionType.ASSIGN_PRIORITY,
                patient_id=pid,
                priority_level=priority,
            )
            obs = env.step(action)

        assert obs.done is True


class TestState:
    """Test state() behavior."""

    def test_state_serializable(self):
        """state() should return serializable data."""
        env = TriageEnvironment()
        env.reset(task_name="basic-triage")
        state = env.state
        d = state.model_dump()
        assert isinstance(d, dict)
        assert "patients" in d
        assert "assignments" in d

    def test_state_tracks_assignments(self):
        """state() should reflect assignments made."""
        env = TriageEnvironment()
        env.reset(task_name="basic-triage")

        action = TriageAction(
            action_type=ActionType.ASSIGN_PRIORITY,
            patient_id="P001",
            priority_level=PriorityLevel.IMMEDIATE,
        )
        env.step(action)

        state = env.state
        assert "P001" in state.assignments
        assert state.assignments["P001"] == "immediate"