tests/test_env.py

"""
Tests for the DevOps RL Environment.
"""

import pytest
from devops_env.env import DevOpsEnv
from scenarios.registry import ScenarioRegistry, Scenario


def _make_test_registry():
    """Create a minimal registry for testing."""
    registry = ScenarioRegistry()
    registry.register(Scenario(
        id="test_missing_pkg",
        level=1,
        description="Test: missing package",
        initial_state={},
        success_condition=lambda out: "successfully installed" in out.lower(),
        hint_commands=["pip install testpkg"],
        error_fingerprint=r"ModuleNotFoundError",
        initial_error_log="ModuleNotFoundError: No module named 'testpkg'",
    ))
    return registry


class TestDevOpsEnv:
    """Tests for the OpenEnv-style environment."""

    def test_reset_returns_observation(self):
        registry = _make_test_registry()
        env = DevOpsEnv(scenario_registry=registry, max_steps=5)
        obs, info = env.reset(options={"scenario_id": "test_missing_pkg"})

        assert "error_log" in obs
        assert "command_history" in obs
        assert "step_count" in obs
        assert "scenario_id" in obs
        assert "error_type" in obs
        assert "error_confidence" in obs
        assert "is_terminal" in obs
        assert "solved" in obs
        assert obs["scenario_id"] == "test_missing_pkg"
        assert obs["step_count"] == 0
        env.close()

    def test_step_returns_correct_tuple(self):
        registry = _make_test_registry()
        env = DevOpsEnv(scenario_registry=registry, max_steps=5)
        obs, _ = env.reset(options={"scenario_id": "test_missing_pkg"})

        result = env.step("echo hello")
        assert len(result) == 5  # obs, reward, terminated, truncated, info

        obs, reward, terminated, truncated, info = result
        assert isinstance(reward, float)
        assert isinstance(terminated, bool)
        assert isinstance(truncated, bool)
        assert isinstance(info, dict)
        env.close()

    def test_max_steps_truncates(self):
        registry = _make_test_registry()
        env = DevOpsEnv(scenario_registry=registry, max_steps=3)
        env.reset(options={"scenario_id": "test_missing_pkg"})

        truncated = False
        for i in range(3):
            _, _, terminated, truncated, _ = env.step("echo step")
            if terminated or truncated:
                break

        assert truncated or terminated  # Should end by step 3
        env.close()

    def test_error_type_classified(self):
        registry = _make_test_registry()
        env = DevOpsEnv(scenario_registry=registry)
        obs, info = env.reset(options={"scenario_id": "test_missing_pkg"})

        assert obs["error_type"] == "missing_package"
        env.close()

    def test_episode_summary(self):
        registry = _make_test_registry()
        env = DevOpsEnv(scenario_registry=registry, max_steps=2)
        env.reset(options={"scenario_id": "test_missing_pkg"})
        env.step("echo test")

        summary = env.get_episode_summary()
        assert summary["scenario_id"] == "test_missing_pkg"
        assert summary["level"] == 1
        assert len(summary["steps"]) == 1
        env.close()

    def test_cannot_step_after_done(self):
        registry = _make_test_registry()
        env = DevOpsEnv(scenario_registry=registry, max_steps=1)
        env.reset(options={"scenario_id": "test_missing_pkg"})
        env.step("echo done")

        with pytest.raises(RuntimeError):
            env.step("echo again")
        env.close()


class TestScenarioRegistry:
    """Tests for the scenario registry."""

    def test_register_and_get(self):
        registry = ScenarioRegistry()
        scenario = Scenario(
            id="test1", level=1, description="test",
            initial_state={}, success_condition=lambda x: True,
            hint_commands=["echo"], error_fingerprint="test",
        )
        registry.register(scenario)
        assert registry.get("test1").id == "test1"

    def test_duplicate_registration_raises(self):
        registry = ScenarioRegistry()
        scenario = Scenario(
            id="dup", level=1, description="test",
            initial_state={}, success_condition=lambda x: True,
            hint_commands=["echo"], error_fingerprint="test",
        )
        registry.register(scenario)
        with pytest.raises(ValueError):
            registry.register(scenario)

    def test_get_by_level(self):
        registry = ScenarioRegistry()
        registry.register_defaults()
        level1 = registry.get_by_level(1)
        assert len(level1) >= 3
        assert all(s.level == 1 for s in level1)

    def test_get_random(self):
        registry = ScenarioRegistry()
        registry.register_defaults()
        scenario = registry.get_random(level=1)
        assert scenario.level == 1

    def test_register_defaults(self):
        registry = ScenarioRegistry()
        registry.register_defaults()
        assert registry.count >= 9  # 3+ per level