tests/test_environment.py

"""Tests for the HR Productivity Environment."""

import pytest

from hr_env.models import HRAction, HRObservation, HRState
from hr_env.server.environment import HRProductivityEnvironment


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


class TestReset:
    def test_returns_observation(self, env):
        obs = env.reset(seed=42)
        assert isinstance(obs, HRObservation)
        assert not obs.done
        assert obs.current_phase == "scanning"
        assert obs.current_quarter == 1

    def test_deterministic_with_seed(self, env):
        obs1 = env.reset(seed=42)
        obs2 = env.reset(seed=42)
        assert obs1.data == obs2.data

    def test_has_company_summary(self, env):
        obs = env.reset(seed=42)
        assert "company_summary" in obs.data
        assert "baseline_metrics" in obs.data

    def test_available_actions_in_scanning(self, env):
        obs = env.reset(seed=42)
        assert "query_department" in obs.available_actions
        assert "calculate_metric" in obs.available_actions


class TestStep:
    def test_query_department(self, env):
        env.reset(seed=42)
        action = HRAction(action_type="query_department", department="Engineering")
        obs = env.step(action)
        assert isinstance(obs, HRObservation)
        assert not obs.done
        assert obs.data is not None
        assert "headcount" in obs.data

    def test_invalid_action_in_phase(self, env):
        env.reset(seed=42)
        # Try executing hiring in scanning phase
        action = HRAction(action_type="execute_hiring", department="Engineering", count=5)
        obs = env.step(action)
        assert "Invalid action" in obs.message or "Invalid" in obs.warnings[0]


class TestDeterminism:
    """Full-episode determinism: stochastic turnover and hire generation must be
    seeded from the episode seed, so the same seed yields the same score."""

    def _run_episode(self, seed):
        from demo import run_demo_json
        return run_demo_json(seed=seed, size=300)["final"]["score"]

    def test_full_episode_reproducible(self):
        a = self._run_episode(42)
        b = self._run_episode(42)
        assert a == b, f"non-deterministic episode: {a} != {b}"

    def test_different_seeds_differ(self):
        # Different seeds should generally give different trajectories.
        assert self._run_episode(42) != self._run_episode(789)

    def test_calculate_metric(self, env):
        env.reset(seed=42)
        action = HRAction(action_type="calculate_metric", metric_name="hcva")
        obs = env.step(action)
        assert obs.data is not None
        assert "hcva" in obs.data

    def test_phase_advancement(self, env):
        env.reset(seed=42)
        # Take minimum scanning actions
        env.step(HRAction(action_type="query_department", department="Engineering"))
        env.step(HRAction(action_type="review_financials"))
        # Advance to planning
        obs = env.step(HRAction(action_type="advance_phase"))
        assert obs.current_phase == "planning"

    def test_cannot_advance_without_min_actions(self, env):
        env.reset(seed=42)
        # Try to advance immediately
        obs = env.step(HRAction(action_type="advance_phase"))
        # Should either be invalid or show warning
        assert "Cannot advance" in obs.message or obs.warnings


class TestQuarterAdvancement:
    def test_full_quarter_cycle(self, env):
        env.reset(seed=42)
        # Scanning
        env.step(HRAction(action_type="query_department", department="Engineering"))
        env.step(HRAction(action_type="review_financials"))
        env.step(HRAction(action_type="advance_phase"))
        # Planning
        env.step(HRAction(action_type="set_hiring_target", department="Engineering", count=5))
        env.step(HRAction(action_type="advance_phase"))
        # Producing
        env.step(HRAction(action_type="execute_hiring", department="Engineering", count=3))
        env.step(HRAction(action_type="advance_phase"))
        # Controlling
        env.step(HRAction(action_type="submit_report"))
        obs = env.step(HRAction(action_type="advance_quarter"))
        # Should advance to Q2
        assert obs.current_quarter == 2
        assert obs.current_phase == "scanning"

    def test_cannot_advance_quarter_outside_controlling(self, env):
        env.reset(seed=42)
        obs = env.step(HRAction(action_type="advance_quarter"))
        assert "controlling" in obs.message.lower() or obs.warnings


class TestState:
    def test_state_after_reset(self, env):
        env.reset(seed=42)
        state = env.state
        assert isinstance(state, HRState)
        assert state.current_quarter == 1
        assert state.current_phase == "scanning"

    def test_state_updates_after_steps(self, env):
        env.reset(seed=42)
        env.step(HRAction(action_type="query_department", department="Engineering"))
        state = env.state
        assert state.step_count == 1


class TestEpisodeCompletion:
    def test_completes_after_6_quarters(self, env):
        """Run through 6 quarters with minimal actions."""
        env.reset(seed=42)

        for q in range(6):
            # Scanning (2 min actions)
            env.step(HRAction(action_type="query_department", department="Engineering"))
            env.step(HRAction(action_type="review_financials"))
            env.step(HRAction(action_type="advance_phase"))
            # Planning (1 min action)
            env.step(HRAction(action_type="set_hiring_target", department="Engineering", count=1))
            env.step(HRAction(action_type="advance_phase"))
            # Producing (1 min action)
            env.step(HRAction(action_type="execute_hiring", department="Engineering", count=1))
            env.step(HRAction(action_type="advance_phase"))
            # Controlling (1 min action)
            env.step(HRAction(action_type="submit_report"))
            obs = env.step(HRAction(action_type="advance_quarter"))

            if q < 5:
                assert not obs.done, f"Episode ended prematurely at Q{q + 1}"
            else:
                assert obs.done, "Episode should be done after Q6"
                assert obs.reward is not None
                assert 0 <= obs.reward <= 1.0

    def test_final_score_in_range(self, env):
        """Verify final score is bounded [0, 1]."""
        env.reset(seed=42)
        for _ in range(6):
            env.step(HRAction(action_type="query_department", department="Sales"))
            env.step(HRAction(action_type="calculate_metric", metric_name="hcva"))
            env.step(HRAction(action_type="advance_phase"))
            env.step(HRAction(action_type="set_training_budget", department="Engineering", amount=10000))
            env.step(HRAction(action_type="advance_phase"))
            env.step(HRAction(action_type="execute_training", department="Engineering", amount=10))
            env.step(HRAction(action_type="advance_phase"))
            env.step(HRAction(action_type="submit_report"))
            obs = env.step(HRAction(action_type="advance_quarter"))
        assert 0 <= obs.reward <= 1.0