tests/test_environment.py

"""
Unit and integration tests for the Email Triage environment.

Tests cover: reset/step/state interface, reward computation, grader
determinism, episode termination, and end-to-end workflow.
"""

import pytest

from src.environment import EmailTriageEnv, TASK_CONFIG
from src.models import Action, Email, Observation, Reward, AVAILABLE_FOLDERS


# ── Reset Tests ──────────────────────────────────────────────────────────────

class TestReset:
    def test_reset_returns_observation(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        assert isinstance(obs, Observation)
        assert not obs.done
        assert obs.current_step == 0

    def test_reset_populates_inbox(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        assert len(obs.inbox_emails) == 5

    def test_reset_medium_task(self):
        env = EmailTriageEnv(task_id="multi_folder_triage", seed=42)
        obs = env.reset()
        assert len(obs.inbox_emails) == 15

    def test_reset_hard_task(self):
        env = EmailTriageEnv(task_id="advanced_triage_with_urgency", seed=42)
        obs = env.reset()
        assert len(obs.inbox_emails) == 30

    def test_reset_deterministic(self):
        """Same seed produces identical emails."""
        env1 = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs1 = env1.reset()

        env2 = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs2 = env2.reset()

        for e1, e2 in zip(obs1.inbox_emails, obs2.inbox_emails):
            assert e1.subject == e2.subject
            assert e1.sender == e2.sender
            assert e1.ground_truth_folder == e2.ground_truth_folder

    def test_reset_different_seeds(self):
        """Different seeds produce different emails."""
        env1 = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs1 = env1.reset()

        env2 = EmailTriageEnv(task_id="basic_triage", seed=99)
        obs2 = env2.reset()

        subjects1 = {e.subject for e in obs1.inbox_emails}
        subjects2 = {e.subject for e in obs2.inbox_emails}
        # Very unlikely to be identical with different seeds
        assert subjects1 != subjects2

    def test_reset_clears_state(self):
        """Reset clears previous episode state."""
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        action = Action(action_type="move", email_id=0, target_folder="work")
        env.step(action)

        obs2 = env.reset()
        assert obs2.current_step == 0
        assert not obs2.done
        assert len(env.history) == 0

    def test_invalid_task_id(self):
        with pytest.raises(ValueError, match="Unknown task_id"):
            EmailTriageEnv(task_id="nonexistent_task")


# ── Step Tests ───────────────────────────────────────────────────────────────

class TestStep:
    def test_step_returns_tuple(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        email_id = obs.inbox_emails[0].id
        action = Action(action_type="move", email_id=email_id, target_folder="work")

        result = env.step(action)
        assert len(result) == 4
        obs, reward, done, info = result
        assert isinstance(obs, Observation)
        assert isinstance(reward, Reward)
        assert isinstance(done, bool)
        assert isinstance(info, dict)

    def test_step_reward_in_range(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        action = Action(
            action_type="move",
            email_id=obs.inbox_emails[0].id,
            target_folder="work",
        )
        _, reward, _, _ = env.step(action)
        assert 0.0 <= reward.value <= 1.0

    def test_step_removes_email_from_inbox(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        initial_count = len(obs.inbox_emails)
        action = Action(
            action_type="move",
            email_id=obs.inbox_emails[0].id,
            target_folder="work",
        )
        obs, _, _, _ = env.step(action)
        assert len(obs.inbox_emails) == initial_count - 1

    def test_step_increments_counter(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        assert obs.current_step == 0

        action = Action(
            action_type="move",
            email_id=obs.inbox_emails[0].id,
            target_folder="work",
        )
        obs, _, _, _ = env.step(action)
        assert obs.current_step == 1

    def test_step_invalid_email_id(self):
        """Step with invalid email_id returns zero reward."""
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        env.reset()
        action = Action(action_type="move", email_id=999, target_folder="work")
        obs, reward, done, info = env.step(action)
        assert reward.value == 0.0
        assert "error" in info

    def test_step_after_done_raises(self):
        """Stepping after done raises RuntimeError."""
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()

        # Process all emails
        while not obs.done:
            if not obs.inbox_emails:
                break
            action = Action(
                action_type="move",
                email_id=obs.inbox_emails[0].id,
                target_folder="work",
            )
            obs, _, _, _ = env.step(action)

        assert obs.done
        with pytest.raises(RuntimeError, match="Episode is done"):
            env.step(Action(action_type="move", email_id=0, target_folder="work"))

    def test_flag_does_not_remove(self):
        """Flagging an email should NOT remove it from inbox."""
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        count_before = len(obs.inbox_emails)
        action = Action(action_type="flag", email_id=obs.inbox_emails[0].id)
        obs, _, _, _ = env.step(action)
        assert len(obs.inbox_emails) == count_before

    def test_mark_spam_removes_email(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        count_before = len(obs.inbox_emails)
        action = Action(action_type="mark_spam", email_id=obs.inbox_emails[0].id)
        obs, _, _, _ = env.step(action)
        assert len(obs.inbox_emails) == count_before - 1


# ── State Tests ──────────────────────────────────────────────────────────────

class TestState:
    def test_state_matches_observation(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()
        state = env.state()
        assert state.current_step == obs.current_step
        assert len(state.inbox_emails) == len(obs.inbox_emails)

    def test_state_does_not_advance(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        env.reset()
        s1 = env.state()
        s2 = env.state()
        assert s1.current_step == s2.current_step


# ── Episode Termination ─────────────────────────────────────────────────────

class TestTermination:
    def test_done_when_inbox_empty(self):
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()

        while obs.inbox_emails:
            action = Action(
                action_type="move",
                email_id=obs.inbox_emails[0].id,
                target_folder="work",
            )
            obs, _, done, _ = env.step(action)
            if done:
                break

        assert obs.done

    def test_done_at_max_steps(self):
        """Episode ends at max_steps even if inbox has emails."""
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()

        # Flag repeatedly (doesn't remove emails) until max steps
        for _ in range(env.max_steps):
            if obs.done:
                break
            action = Action(action_type="flag", email_id=obs.inbox_emails[0].id)
            obs, _, _, _ = env.step(action)

        assert obs.done


# ── Grader Determinism ───────────────────────────────────────────────────────

class TestGraderDeterminism:
    def _run_episode(self, task_id: str, seed: int) -> float:
        """Run a full episode and return the grader score."""
        from src.grader import grade_task_basic, grade_task_medium, grade_task_hard

        graders = {
            "basic_triage": grade_task_basic,
            "multi_folder_triage": grade_task_medium,
            "advanced_triage_with_urgency": grade_task_hard,
        }

        env = EmailTriageEnv(task_id=task_id, seed=seed)
        obs = env.reset()

        while not obs.done and obs.inbox_emails:
            email = obs.inbox_emails[0]
            # Use ground truth for a "perfect" agent
            action = Action(
                action_type="move",
                email_id=email.id,
                target_folder=email.ground_truth_folder,
                reason="Ground truth classification",
            )
            obs, _, _, _ = env.step(action)

        return graders[task_id](env.history)

    def test_basic_grader_deterministic(self):
        s1 = self._run_episode("basic_triage", seed=42)
        s2 = self._run_episode("basic_triage", seed=42)
        assert s1 == s2

    def test_medium_grader_deterministic(self):
        s1 = self._run_episode("multi_folder_triage", seed=42)
        s2 = self._run_episode("multi_folder_triage", seed=42)
        assert s1 == s2

    def test_hard_grader_deterministic(self):
        s1 = self._run_episode("advanced_triage_with_urgency", seed=42)
        s2 = self._run_episode("advanced_triage_with_urgency", seed=42)
        assert s1 == s2

    def test_perfect_agent_scores_high(self):
        """An agent using ground truth should score near 1.0."""
        score = self._run_episode("basic_triage", seed=42)
        assert score >= 0.8, f"Perfect agent scored {score}, expected >= 0.8"

    def test_grader_returns_float_in_range(self):
        for task_id in TASK_CONFIG:
            score = self._run_episode(task_id, seed=42)
            assert isinstance(score, float)
            assert 0.0 < score < 1.0, f"{task_id}: score {score} out of strict (0,1) range"


# ── Integration Test ─────────────────────────────────────────────────────────

class TestIntegration:
    def test_full_episode_workflow(self):
        """End-to-end: reset, step through all emails, check grading."""
        env = EmailTriageEnv(task_id="basic_triage", seed=42)
        obs = env.reset()

        rewards = []
        while not obs.done and obs.inbox_emails:
            email = obs.inbox_emails[0]
            action = Action(
                action_type="move",
                email_id=email.id,
                target_folder=email.ground_truth_folder,
            )
            obs, reward, done, info = env.step(action)
            rewards.append(reward.value)

        assert obs.done
        assert len(rewards) == 5
        assert all(0.0 <= r <= 1.0 for r in rewards)
        assert len(env.history) == 5

    def test_all_tasks_runnable(self):
        """All 3 tasks can be reset and stepped without errors."""
        for task_id in TASK_CONFIG:
            env = EmailTriageEnv(task_id=task_id, seed=42)
            obs = env.reset()
            assert not obs.done
            assert len(obs.inbox_emails) > 0

            action = Action(
                action_type="move",
                email_id=obs.inbox_emails[0].id,
                target_folder="work",
            )
            obs, reward, done, info = env.step(action)
            assert reward.value >= 0.0