| """Unit tests for Reward_Function.
|
|
|
| Tests exact R1/R2/R3/R4 values and penalty application.
|
| """
|
| from __future__ import annotations
|
|
|
| import pytest
|
|
|
| from sentinel.models import (
|
| Action,
|
| FailureType,
|
| IncidentState,
|
| RewardBreakdown,
|
| RewardWeights,
|
| ServiceMetrics,
|
| Trajectory,
|
| TrajectoryStep,
|
| )
|
| from sentinel.reward import Reward_Function
|
| from sentinel.world_state import ALL_SERVICES, NexaStackWorldState
|
|
|
|
|
|
|
|
|
|
|
| SLA_THRESHOLD = 50
|
|
|
| WEIGHTS = RewardWeights(
|
| r1_root_cause=0.35,
|
| r2_mttr=0.30,
|
| r3_recovery_quality=0.25,
|
| r4_blast_radius=0.10,
|
| )
|
|
|
|
|
| @pytest.fixture
|
| def rf() -> Reward_Function:
|
| return Reward_Function(weights=WEIGHTS, sla_breach_threshold=SLA_THRESHOLD)
|
|
|
|
|
| def _make_incident_state(
|
| root_cause_service: str = "redis-cache",
|
| failure_type: FailureType = FailureType.memory_leak,
|
| peak_blast_radius: set[str] | None = None,
|
| current_blast_radius: set[str] | None = None,
|
| ) -> IncidentState:
|
| if peak_blast_radius is None:
|
| peak_blast_radius = set()
|
| if current_blast_radius is None:
|
| current_blast_radius = set()
|
| return IncidentState(
|
| template_id="E1",
|
| root_cause_service=root_cause_service,
|
| failure_type=failure_type,
|
| ground_truth_signals=[],
|
| red_herring_signals=[],
|
| affected_services={},
|
| peak_blast_radius=peak_blast_radius,
|
| current_blast_radius=current_blast_radius,
|
| timeline=[],
|
| attempted_remediations=[],
|
| active_hypotheses=[],
|
| resolved=False,
|
| step_injected=0,
|
| )
|
|
|
|
|
| def _make_trajectory(
|
| mttr: int,
|
| identified_root_cause: str = "",
|
| identified_failure_type: str = "",
|
| ) -> Trajectory:
|
| action = Action(
|
| agent="holmes",
|
| category="meta",
|
| name="CloseIncident",
|
| params={"resolution_summary": "done"},
|
| )
|
| step = TrajectoryStep(
|
| observation={},
|
| action=action,
|
| reward=0.0,
|
| terminated=True,
|
| truncated=False,
|
| info={
|
| "identified_root_cause": identified_root_cause,
|
| "identified_failure_type": identified_failure_type,
|
| },
|
| )
|
| return Trajectory(
|
| episode_id="test-ep",
|
| incident_template_id="E1",
|
| steps=[step],
|
| final_reward=RewardBreakdown(r1=0.0, r2=0.0, r3=0.0, r4=0.0, penalties=0.0, total=0.0),
|
| mttr=mttr,
|
| )
|
|
|
|
|
| def _baseline_world_state() -> NexaStackWorldState:
|
| ws = NexaStackWorldState()
|
| ws.restore_baseline()
|
| return ws
|
|
|
|
|
|
|
|
|
|
|
|
|
| class TestR1RootCauseAccuracy:
|
| def test_r1_correct_service_and_type_returns_1_0(self, rf):
|
| """Correct service + correct failure type → R1 = 1.0"""
|
| incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| failure_type=FailureType.memory_leak,
|
| )
|
| r1 = rf._r1_root_cause_accuracy(incident, "redis-cache", "memory_leak")
|
| assert r1 == 1.0
|
|
|
| def test_r1_correct_service_wrong_type_returns_0_5(self, rf):
|
| """Correct service + wrong failure type → R1 = 0.5"""
|
| incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| failure_type=FailureType.memory_leak,
|
| )
|
| r1 = rf._r1_root_cause_accuracy(incident, "redis-cache", "cpu_spike")
|
| assert r1 == 0.5
|
|
|
| def test_r1_wrong_service_returns_0_0(self, rf):
|
| """Wrong service → R1 = 0.0 regardless of failure type"""
|
| incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| failure_type=FailureType.memory_leak,
|
| )
|
| r1 = rf._r1_root_cause_accuracy(incident, "postgres-primary", "memory_leak")
|
| assert r1 == 0.0
|
|
|
| def test_r1_wrong_service_wrong_type_returns_0_0(self, rf):
|
| """Wrong service + wrong type → R1 = 0.0"""
|
| incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| failure_type=FailureType.memory_leak,
|
| )
|
| r1 = rf._r1_root_cause_accuracy(incident, "kafka-broker", "cpu_spike")
|
| assert r1 == 0.0
|
|
|
| def test_r1_empty_identified_returns_0_0(self, rf):
|
| """Empty identified root cause → R1 = 0.0"""
|
| incident = _make_incident_state(root_cause_service="redis-cache")
|
| r1 = rf._r1_root_cause_accuracy(incident, "", "")
|
| assert r1 == 0.0
|
|
|
|
|
|
|
|
|
|
|
|
|
| class TestR2Mttr:
|
| def test_r2_mttr_zero_returns_1_1(self, rf):
|
| """mttr=0 → R2 = 1.0 / (1.0 + 0/50) + 0.1 = 1.1"""
|
| r2 = rf._r2_mttr(0)
|
| assert abs(r2 - 1.1) < 1e-9, f"Expected 1.1, got {r2}"
|
|
|
| def test_r2_mttr_at_sla_threshold_returns_0_5(self, rf):
|
| """mttr=sla_threshold → R2 = 1.0 / (1.0 + 1.0) = 0.5 (no bonus)"""
|
| r2 = rf._r2_mttr(SLA_THRESHOLD)
|
|
|
| expected = 1.0 / (1.0 + SLA_THRESHOLD / SLA_THRESHOLD)
|
| assert abs(r2 - expected) < 1e-9, f"Expected {expected}, got {r2}"
|
|
|
| def test_r2_mttr_double_threshold_returns_one_third(self, rf):
|
| """mttr=2*sla_threshold → R2 = 1.0 / (1.0 + 2.0) = 1/3"""
|
| mttr = 2 * SLA_THRESHOLD
|
| r2 = rf._r2_mttr(mttr)
|
| expected = 1.0 / (1.0 + mttr / SLA_THRESHOLD)
|
| assert abs(r2 - expected) < 1e-9, f"Expected {expected:.6f}, got {r2:.6f}"
|
|
|
| def test_r2_pre_sla_bonus_applied_when_mttr_below_threshold(self, rf):
|
| """mttr < sla_threshold → bonus of 0.1 is added"""
|
| mttr = SLA_THRESHOLD - 1
|
| r2 = rf._r2_mttr(mttr)
|
| base = 1.0 / (1.0 + mttr / SLA_THRESHOLD)
|
| expected = base + 0.1
|
| assert abs(r2 - expected) < 1e-9, f"Expected {expected:.6f}, got {r2:.6f}"
|
|
|
| def test_r2_no_bonus_at_threshold(self, rf):
|
| """mttr == sla_threshold → no bonus"""
|
| r2_at = rf._r2_mttr(SLA_THRESHOLD)
|
| r2_below = rf._r2_mttr(SLA_THRESHOLD - 1)
|
|
|
| assert r2_below > r2_at
|
|
|
| def test_r2_clamped_to_1_1(self, rf):
|
| """R2 is clamped to max 1.1"""
|
| r2 = rf._r2_mttr(0)
|
| assert r2 <= 1.1
|
|
|
| def test_r2_clamped_to_0_0(self, rf):
|
| """R2 is clamped to min 0.0"""
|
| r2 = rf._r2_mttr(10_000_000)
|
| assert r2 >= 0.0
|
|
|
|
|
|
|
|
|
|
|
|
|
| class TestR3RecoveryQuality:
|
| def test_r3_all_services_at_baseline_returns_1_0(self, rf):
|
| """All services at baseline → R3 = 1.0"""
|
| ws = _baseline_world_state()
|
| r3 = rf._r3_recovery_quality(ws)
|
| assert abs(r3 - 1.0) < 1e-9, f"Expected 1.0, got {r3}"
|
|
|
| def test_r3_all_services_degraded_returns_0_0(self, rf):
|
| """All services severely degraded → R3 = 0.0"""
|
| ws = _baseline_world_state()
|
|
|
| for svc in ALL_SERVICES:
|
| ws.services[svc] = ServiceMetrics(
|
| cpu=1.0,
|
| memory=1.0,
|
| latency_ms=1000.0,
|
| error_rate=1.0,
|
| saturation=1.0,
|
| availability=False,
|
| )
|
| r3 = rf._r3_recovery_quality(ws)
|
| assert abs(r3 - 0.0) < 1e-9, f"Expected 0.0, got {r3}"
|
|
|
| def test_r3_half_services_recovered_returns_0_5(self, rf):
|
| """Half services at baseline, half degraded → R3 = 0.5"""
|
| ws = _baseline_world_state()
|
| half = len(ALL_SERVICES) // 2
|
| for svc in ALL_SERVICES[:half]:
|
| ws.services[svc] = ServiceMetrics(
|
| cpu=1.0,
|
| memory=1.0,
|
| latency_ms=1000.0,
|
| error_rate=1.0,
|
| saturation=1.0,
|
| availability=False,
|
| )
|
| r3 = rf._r3_recovery_quality(ws)
|
| assert abs(r3 - 0.5) < 1e-9, f"Expected 0.5, got {r3}"
|
|
|
| def test_r3_empty_services_returns_1_0(self, rf):
|
| """No services → R3 = 1.0 (vacuously true)"""
|
| ws = _baseline_world_state()
|
| ws.services = {}
|
| r3 = rf._r3_recovery_quality(ws)
|
| assert abs(r3 - 1.0) < 1e-9, f"Expected 1.0, got {r3}"
|
|
|
|
|
|
|
|
|
|
|
|
|
| class TestR4BlastRadius:
|
| def test_r4_current_equals_peak_returns_0_0(self, rf):
|
| """current_br == peak_br → R4 = 0.0"""
|
| services = {"svc-a", "svc-b", "svc-c"}
|
| incident = _make_incident_state(
|
| peak_blast_radius=services,
|
| current_blast_radius=services,
|
| )
|
| r4 = rf._r4_blast_radius(incident)
|
| assert abs(r4 - 0.0) < 1e-9, f"Expected 0.0, got {r4}"
|
|
|
| def test_r4_current_zero_returns_1_0(self, rf):
|
| """current_br = 0, peak_br > 0 → R4 = 1.0"""
|
| incident = _make_incident_state(
|
| peak_blast_radius={"svc-a", "svc-b"},
|
| current_blast_radius=set(),
|
| )
|
| r4 = rf._r4_blast_radius(incident)
|
| assert abs(r4 - 1.0) < 1e-9, f"Expected 1.0, got {r4}"
|
|
|
| def test_r4_peak_zero_returns_1_0(self, rf):
|
| """peak_br = 0 → R4 = 1.0 (no blast occurred)"""
|
| incident = _make_incident_state(
|
| peak_blast_radius=set(),
|
| current_blast_radius=set(),
|
| )
|
| r4 = rf._r4_blast_radius(incident)
|
| assert abs(r4 - 1.0) < 1e-9, f"Expected 1.0, got {r4}"
|
|
|
| def test_r4_partial_reduction(self, rf):
|
| """current_br = half of peak_br → R4 = 0.5"""
|
| peak = {"a", "b", "c", "d"}
|
| current = {"a", "b"}
|
| incident = _make_incident_state(
|
| peak_blast_radius=peak,
|
| current_blast_radius=current,
|
| )
|
| r4 = rf._r4_blast_radius(incident)
|
| expected = 1.0 - (2 / 4)
|
| assert abs(r4 - expected) < 1e-9, f"Expected {expected}, got {r4}"
|
|
|
|
|
|
|
|
|
|
|
|
|
| class TestPenalties:
|
| def test_blast_radius_expansion_penalty(self, rf):
|
| """Blast radius expansion → step reward includes -1.0 base penalty."""
|
| pre_br = {"svc-a"}
|
| post_br = {"svc-a", "svc-b"}
|
|
|
|
|
|
|
| pre_incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| current_blast_radius=pre_br,
|
| )
|
| ws = _baseline_world_state()
|
| ws.incident_state = _make_incident_state(current_blast_radius=post_br)
|
|
|
| action = Action(
|
| agent="forge",
|
| category="remediation",
|
| name="ScaleService",
|
| params={"service": "api-gateway", "replicas": 3},
|
| )
|
|
|
| step_reward = rf.compute_step_reward(action, ws, pre_incident)
|
|
|
|
|
|
|
| assert step_reward == -1.0, f"Expected -1.0, got {step_reward}"
|
|
|
| def test_no_blast_radius_penalty_when_br_unchanged(self, rf):
|
| """No blast radius change → no -1.0 penalty"""
|
| br = {"svc-a", "svc-b"}
|
|
|
| pre_incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| current_blast_radius=br,
|
| )
|
| ws = _baseline_world_state()
|
| ws.incident_state = _make_incident_state(current_blast_radius=br)
|
|
|
| action = Action(
|
| agent="forge",
|
| category="remediation",
|
| name="ScaleService",
|
| params={"service": "api-gateway", "replicas": 3},
|
| )
|
|
|
| step_reward = rf.compute_step_reward(action, ws, pre_incident)
|
| assert step_reward == 0.0, f"Expected 0.0, got {step_reward}"
|
|
|
| def test_restart_healthy_service_penalty(self, rf):
|
| """RestartService on healthy non-root service → step reward = -1.0"""
|
| ws = _baseline_world_state()
|
|
|
| assert ws.services["api-gateway"].availability is True
|
|
|
| pre_incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| current_blast_radius=set(),
|
| )
|
| ws.incident_state = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| current_blast_radius=set(),
|
| )
|
|
|
| action = Action(
|
| agent="forge",
|
| category="remediation",
|
| name="RestartService",
|
| params={"service": "api-gateway"},
|
| )
|
|
|
| step_reward = rf.compute_step_reward(action, ws, pre_incident)
|
| assert step_reward == -1.0, f"Expected -1.0, got {step_reward}"
|
|
|
| def test_restart_unhealthy_service_no_penalty(self, rf):
|
| """RestartService on unhealthy root cause service → +0.25 remediation shaping"""
|
| ws = _baseline_world_state()
|
|
|
| ws.services["redis-cache"] = ServiceMetrics(
|
| cpu=1.0, memory=1.0, latency_ms=1000.0,
|
| error_rate=1.0, saturation=1.0, availability=False,
|
| )
|
|
|
| pre_incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| current_blast_radius={"redis-cache"},
|
| )
|
| ws.incident_state = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| current_blast_radius={"redis-cache"},
|
| )
|
|
|
| action = Action(
|
| agent="forge",
|
| category="remediation",
|
| name="RestartService",
|
| params={"service": "redis-cache"},
|
| )
|
|
|
| step_reward = rf.compute_step_reward(action, ws, pre_incident)
|
|
|
| assert step_reward == 0.25, f"Expected 0.25, got {step_reward}"
|
|
|
| def test_late_resolution_penalty_in_episode_reward(self, rf):
|
| """mttr > 2 * sla_threshold → episode penalties = -0.5"""
|
| ws = _baseline_world_state()
|
| incident = _make_incident_state()
|
| trajectory = _make_trajectory(mttr=2 * SLA_THRESHOLD + 1)
|
|
|
| result = rf.compute_episode_reward(trajectory, ws, incident)
|
| assert result.penalties == -0.5, f"Expected -0.5, got {result.penalties}"
|
|
|
| def test_no_late_penalty_when_resolved_on_time(self, rf):
|
| """mttr <= 2 * sla_threshold → episode penalties = 0.0"""
|
| ws = _baseline_world_state()
|
| incident = _make_incident_state()
|
| trajectory = _make_trajectory(mttr=2 * SLA_THRESHOLD)
|
|
|
| result = rf.compute_episode_reward(trajectory, ws, incident)
|
| assert result.penalties == 0.0, f"Expected 0.0, got {result.penalties}"
|
|
|
| def test_both_penalties_accumulate(self, rf):
|
| """Blast radius expansion + healthy restart on non-root service → step reward = -2.0"""
|
| pre_br = {"svc-a"}
|
| post_br = {"svc-a", "svc-b"}
|
|
|
| pre_incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| current_blast_radius=pre_br,
|
| )
|
| ws = _baseline_world_state()
|
| ws.incident_state = _make_incident_state(current_blast_radius=post_br)
|
|
|
|
|
| action = Action(
|
| agent="forge",
|
| category="remediation",
|
| name="RestartService",
|
| params={"service": "api-gateway"},
|
| )
|
|
|
| step_reward = rf.compute_step_reward(action, ws, pre_incident)
|
| assert step_reward == -2.0, f"Expected -2.0, got {step_reward}"
|
|
|
|
|
|
|
|
|
|
|
|
|
| class TestComputeEpisodeReward:
|
| def test_total_equals_weighted_sum(self, rf):
|
| """total == 0.35*R1 + 0.30*R2 + 0.25*R3 + 0.10*R4 + penalties"""
|
| ws = _baseline_world_state()
|
| incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| failure_type=FailureType.memory_leak,
|
| peak_blast_radius={"a", "b", "c", "d"},
|
| current_blast_radius={"a", "b"},
|
| )
|
| trajectory = _make_trajectory(
|
| mttr=10,
|
| identified_root_cause="redis-cache",
|
| identified_failure_type="memory_leak",
|
| )
|
|
|
| result = rf.compute_episode_reward(trajectory, ws, incident)
|
|
|
| expected_total = (
|
| 0.35 * result.r1
|
| + 0.30 * result.r2
|
| + 0.25 * result.r3
|
| + 0.10 * result.r4
|
| + result.penalties
|
| )
|
| assert abs(result.total - expected_total) < 1e-9
|
|
|
| def test_perfect_episode_reward(self, rf):
|
| """Perfect identification + fast MTTR + all recovered + no blast → high reward"""
|
| ws = _baseline_world_state()
|
| incident = _make_incident_state(
|
| root_cause_service="redis-cache",
|
| failure_type=FailureType.memory_leak,
|
| peak_blast_radius=set(),
|
| current_blast_radius=set(),
|
| )
|
| trajectory = _make_trajectory(
|
| mttr=0,
|
| identified_root_cause="redis-cache",
|
| identified_failure_type="memory_leak",
|
| )
|
|
|
| result = rf.compute_episode_reward(trajectory, ws, incident)
|
|
|
|
|
| assert result.r1 == 1.0
|
| assert abs(result.r2 - 1.1) < 1e-9
|
| assert abs(result.r3 - 1.0) < 1e-9
|
| assert abs(result.r4 - 1.0) < 1e-9
|
| assert result.penalties == 0.0
|
|
|
| expected_total = 0.35 * 1.0 + 0.30 * 1.1 + 0.25 * 1.0 + 0.10 * 1.0
|
| assert abs(result.total - expected_total) < 1e-9
|
|
|
