tests/test_rewards.py

"""
Tests for the Reward Engine — 100% coverage of all reward signals.
"""

import pytest
from rewards.engine import RewardEngine
from executor.docker_executor import ExecutionResult
from scenarios.registry import Scenario


def _make_scenario(**kwargs):
    """Create a test scenario with sensible defaults."""
    defaults = dict(
        id="test_scenario",
        level=1,
        description="Test scenario",
        initial_state={},
        success_condition=lambda output: "success" in output.lower(),
        hint_commands=["pip install flask"],
        error_fingerprint=r"ModuleNotFoundError",
    )
    defaults.update(kwargs)
    return Scenario(**defaults)


def _make_result(**kwargs):
    """Create a test ExecutionResult with sensible defaults."""
    defaults = dict(stdout="", stderr="", exit_code=0, timed_out=False, blocked=False, block_reason="")
    defaults.update(kwargs)
    return ExecutionResult(**defaults)


class TestRewardSuccess:
    """Test the success reward signal."""

    def test_success_gives_positive_reward(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(stdout="Successfully installed flask\nSuccess")
        total, breakdown = engine.compute_reward(
            action="pip install flask",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["pip install flask"],
            prev_error_log="ModuleNotFoundError",
            curr_error_log="Successfully installed flask",
        )
        assert breakdown.get("success", 0) == 10.0
        assert total > 0

    def test_no_success_gives_no_success_reward(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(stdout="some output", stderr="still broken")
        total, breakdown = engine.compute_reward(
            action="ls",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["ls"],
            prev_error_log="error",
            curr_error_log="still broken",
        )
        assert "success" not in breakdown


class TestRewardCorrectCommand:
    """Test the correct_command reward signal."""

    def test_hint_command_gets_bonus(self):
        engine = RewardEngine()
        scenario = _make_scenario(hint_commands=["pip install flask"])
        result = _make_result(stdout="installed")
        total, breakdown = engine.compute_reward(
            action="pip install flask",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["pip install flask"],
            prev_error_log="error",
            curr_error_log="installed",
        )
        assert breakdown.get("correct_command", 0) == 1.5

    def test_wrong_command_no_bonus(self):
        engine = RewardEngine()
        scenario = _make_scenario(hint_commands=["pip install flask"])
        result = _make_result(stdout="output")
        total, breakdown = engine.compute_reward(
            action="apt-get install python",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["apt-get install python"],
            prev_error_log="error",
            curr_error_log="output",
        )
        assert "correct_command" not in breakdown


class TestRewardProgress:
    """Test the progress and no_progress signals."""

    def test_progress_when_error_changes(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(stdout="new output")
        total, breakdown = engine.compute_reward(
            action="pip install flask",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["pip install flask"],
            prev_error_log="ModuleNotFoundError: No module named 'flask'",
            curr_error_log="installed successfully",
        )
        assert breakdown.get("progress", 0) == 1.0

    def test_no_progress_when_identical_logs(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(stdout="same error")
        same_log = "ModuleNotFoundError: No module named 'flask'"
        total, breakdown = engine.compute_reward(
            action="echo hello",
            result=result,
            scenario=scenario,
            step_count=2,
            command_history=["echo hello"],
            prev_error_log=same_log,
            curr_error_log=same_log,
        )
        assert breakdown.get("no_progress", 0) == -1.0


class TestRewardEfficiency:
    """Test the efficiency_bonus signal."""

    def test_efficiency_bonus_when_solved_fast(self):
        engine = RewardEngine()
        scenario = _make_scenario(hint_commands=["pip install flask"])
        result = _make_result(stdout="Success")
        total, breakdown = engine.compute_reward(
            action="pip install flask",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["pip install flask"],
            prev_error_log="error",
            curr_error_log="Success",
        )
        assert breakdown.get("efficiency_bonus", 0) == 2.0

    def test_no_efficiency_bonus_when_too_many_steps(self):
        engine = RewardEngine()
        scenario = _make_scenario(hint_commands=["pip install flask"])
        result = _make_result(stdout="Success")
        total, breakdown = engine.compute_reward(
            action="pip install flask",
            result=result,
            scenario=scenario,
            step_count=5,
            command_history=["a", "b", "c", "d", "pip install flask"],
            prev_error_log="error",
            curr_error_log="Success",
        )
        assert "efficiency_bonus" not in breakdown


class TestRewardPenalties:
    """Test penalty signals."""

    def test_blocked_command_penalty(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(blocked=True, block_reason="Command 'foo' is not in the whitelist")
        total, breakdown = engine.compute_reward(
            action="foo",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["foo"],
            prev_error_log="error",
            curr_error_log="blocked",
        )
        assert breakdown.get("invalid_command", 0) == -2.0

    def test_dangerous_command_penalty(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(blocked=True, block_reason="Dangerous blocklist pattern matched")
        total, breakdown = engine.compute_reward(
            action="rm -rf /",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["rm -rf /"],
            prev_error_log="error",
            curr_error_log="blocked",
        )
        assert breakdown.get("dangerous_command", 0) == -10.0

    def test_timeout_penalty(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(timed_out=True)
        total, breakdown = engine.compute_reward(
            action="sleep 100",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["sleep 100"],
            prev_error_log="error",
            curr_error_log="timeout",
        )
        assert breakdown.get("timeout", 0) == -5.0

    def test_repeated_command_penalty(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(stdout="output")
        total, breakdown = engine.compute_reward(
            action="pip install flask",
            result=result,
            scenario=scenario,
            step_count=2,
            command_history=["pip install flask", "pip install flask"],
            prev_error_log="error",
            curr_error_log="output",
        )
        assert breakdown.get("repeated_command", 0) == -1.5

    def test_step_cost_always_applied(self):
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(stdout="output")
        total, breakdown = engine.compute_reward(
            action="ls",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["ls"],
            prev_error_log="error",
            curr_error_log="output",
        )
        assert breakdown.get("step_cost", 0) == -0.2


class TestRewardCombinations:
    """Test reward signal combinations."""

    def test_perfect_solve_gives_max_reward(self):
        engine = RewardEngine()
        scenario = _make_scenario(hint_commands=["pip install flask"])
        result = _make_result(stdout="Successfully installed flask. Success")
        total, breakdown = engine.compute_reward(
            action="pip install flask",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["pip install flask"],
            prev_error_log="ModuleNotFoundError",
            curr_error_log="installed",
        )
        # Should get: success(10) + correct_command(1.5) + progress(1) + efficiency(2) + step_cost(-0.2) = 14.3
        assert total > 14.0
        assert "success" in breakdown
        assert "correct_command" in breakdown
        assert "efficiency_bonus" in breakdown

    def test_blocked_command_short_circuits(self):
        """Blocked commands should only get step_cost + the block penalty."""
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(blocked=True, block_reason="not in whitelist")
        total, breakdown = engine.compute_reward(
            action="foo",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["foo"],
            prev_error_log="error",
            curr_error_log="blocked",
        )
        assert len(breakdown) == 2  # step_cost + invalid_command
        assert "success" not in breakdown
        assert "progress" not in breakdown

    def test_timed_out_short_circuits(self):
        """Timed out commands should only get step_cost + timeout penalty."""
        engine = RewardEngine()
        scenario = _make_scenario()
        result = _make_result(timed_out=True)
        total, breakdown = engine.compute_reward(
            action="sleep 999",
            result=result,
            scenario=scenario,
            step_count=1,
            command_history=["sleep 999"],
            prev_error_log="error",
            curr_error_log="timeout",
        )
        assert len(breakdown) == 2  # step_cost + timeout
        assert total == -5.2