tests/test_incident_environment.py

from __future__ import annotations

import unittest

from incident_commander_env.models import IncidentAction
from incident_commander_env.orchestration import run_multi_agent_episode
from incident_commander_env.scenarios import FAMILIES
from incident_commander_env.server.incident_commander_environment import IncidentCommanderEnvironment
from scripts.manual_episode import SOLUTIONS, first_relevant_metric


def step(env: IncidentCommanderEnvironment, tool_name: str, **arguments):
    return env.step(IncidentAction(tool_name=tool_name, arguments=arguments))


class IncidentCommanderEnvironmentTests(unittest.TestCase):
    def test_oracle_solution_scores_high_for_all_families(self):
        for family in FAMILIES:
            with self.subTest(family=family):
                env = IncidentCommanderEnvironment()
                env.reset(scenario_id=family)
                solution = SOLUTIONS[env.scenario.family]
                first = env.scenario.services[0]
                step(env, "query_logs", service=first, query="error warn timeout crash", minutes=20)
                metric_service, metric = first_relevant_metric(env)
                step(env, "inspect_metric", service=metric_service, metric=metric, minutes=30)
                step(env, "read_runbook", topic=solution["topic"])
                step(env, "test_hypothesis", root_cause=solution["root"])
                step(env, "apply_mitigation", mitigation=solution["mitigation"])
                final = step(
                    env,
                    "final_report",
                    root_cause=solution["root"],
                    mitigation=solution["mitigation"],
                    customer_update=(
                        "Customers are impacted by elevated failures. We applied mitigation and "
                        "expect recovery within 15 minutes; next update in 10 minutes."
                    ),
                )
                self.assertTrue(final.done)
                self.assertGreaterEqual(final.reward, 0.85)

    def test_wrong_root_cause_is_capped(self):
        env = IncidentCommanderEnvironment()
        env.reset(scenario_id="cache_saturation")
        final = step(
            env,
            "final_report",
            root_cause="database primary is down",
            mitigation="restart database",
            customer_update="Customers are impacted; mitigation is underway with update in 10 minutes.",
        )
        self.assertLessEqual(final.reward, 0.35)

    def test_unsafe_mitigation_penalized(self):
        env = IncidentCommanderEnvironment()
        env.reset(scenario_id="queue_worker_crash")
        solution = SOLUTIONS[env.scenario.family]
        step(env, "apply_mitigation", mitigation="purge queue and delete messages")
        final = step(
            env,
            "final_report",
            root_cause=solution["root"],
            mitigation="purge queue and delete messages",
            customer_update="Customers are impacted; mitigation is underway with update in 10 minutes.",
        )
        self.assertLess(final.reward, 0.7)
        self.assertIn("unsafe_penalty", final.score_breakdown)

    def test_invalid_tool_does_not_crash(self):
        env = IncidentCommanderEnvironment()
        env.reset(scenario_id="bad_deploy_auth")
        obs = step(env, "delete_everything", service="auth-service")
        self.assertIn("Invalid tool", obs.tool_result)
        self.assertEqual(env.state.invalid_actions, 1)

    def test_excessive_tool_calls_reduce_efficiency(self):
        env = IncidentCommanderEnvironment()
        env.reset(scenario_id="payment_provider_degradation")
        solution = SOLUTIONS[env.scenario.family]
        service = env.scenario.services[0]
        for _ in range(10):
            step(env, "query_logs", service=service, query="error", minutes=5)
        final = step(
            env,
            "final_report",
            root_cause=solution["root"],
            mitigation=solution["mitigation"],
            customer_update=(
                "Customers are impacted. We applied mitigation and expect recovery within "
                "15 minutes; next update in 10 minutes."
            ),
        )
        self.assertLess(final.score_breakdown["safety_efficiency"], 0.10)

    def test_multi_agent_orchestration_solves_representative_incidents(self):
        expected_roots = {
            "payment_provider_degradation": "payment provider degradation",
            "database_connection_leak": "database connection leak",
            "queue_worker_crash": "worker crash",
        }
        for family, expected in expected_roots.items():
            with self.subTest(family=family):
                result = run_multi_agent_episode(scenario_id=family)
                self.assertIn(expected, result.root_cause)
                self.assertGreaterEqual(result.reward, 0.75)
                self.assertEqual(len(result.findings), 3)


if __name__ == "__main__":
    unittest.main()