from __future__ import annotations import contextlib from io import StringIO import json from pathlib import Path import tempfile import unittest from unittest.mock import call, patch from ai_triage.cli import main from ai_triage.monitor import CycleError, CycleResult CONFIG = """ [github] owner = "tidbcloud" repo = "tidb-management-service" token_env = "AI_TRIAGE_GITHUB_TOKEN" [agent] command = ["opencode", "acp"] prompt_path = "prompts/triage_knowledge.md" timeout_seconds = 60 [monitor] poll_interval_seconds = 60 lookback_seconds = 86400 max_pages_per_target = 10 state_path = "state/events.sqlite3" lease_seconds = 600 retry_base_seconds = 10 retry_max_seconds = 60 max_attempts = 3 retention_days = 30 [feishu] webhook_env = "AI_TRIAGE_FEISHU_WEBHOOK" app_id_env = "LARK_APP_ID" app_secret_env = "LARK_APP_SECRET" [oncall] api_key_env = "ONCALL_API_KEY" base_url = "https://oncall.pingcap.net" group_name = "ONCALL" timeout_seconds = 10 cache_ttl_seconds = 60 [routing] api_team_map_path = "routing/api-team-map.toml" [owners] product = "" case = "case@example.com" infrastructure = "" [output] directory = "reports" [limits] max_log_bytes = 100000 max_evidence_chars = 10000 max_agent_output_chars = 10000 max_acp_protocol_bytes = 500000 [[targets]] id = "safe-example" enabled = true workflow = "scenario.yaml" branch = "main" """ class RecordingState: def __init__(self, path: Path) -> None: self.path = path self.closed = False def __enter__(self) -> RecordingState: return self def __exit__(self, *args: object) -> None: self.closed = True class RecordingMonitor: result = CycleResult() instances: list[RecordingMonitor] = [] def __init__(self, *args: object, **kwargs: object) -> None: self.args = args self.kwargs = kwargs self.ran = False self.stop_event = None self.__class__.instances.append(self) def scan_once(self) -> CycleResult: return self.__class__.result def run(self, stop_event: object, *, on_cycle: object = None) -> None: self.ran = True self.stop_event = stop_event self.on_cycle = on_cycle if on_cycle is not None: on_cycle(CycleResult()) getattr(stop_event, "set")() class CLITests(unittest.TestCase): def setUp(self) -> None: RecordingMonitor.instances.clear() RecordingMonitor.result = CycleResult() def fixture(self, directory: str, *, enabled: bool = True) -> Path: root = Path(directory) (root / "prompts").mkdir() (root / "prompts" / "triage_knowledge.md").write_text( "knowledge", encoding="utf-8" ) (root / "routing").mkdir() (root / "routing" / "api-team-map.toml").write_text( """ version = 1 [[routes]] path_pattern = "/api/v2025/cloud/v1beta2/tidbs/*" team = "global_control" """, encoding="utf-8", ) path = root / "config.toml" content = CONFIG.replace("enabled = true", f"enabled = {str(enabled).lower()}") path.write_text(content, encoding="utf-8") return path def test_validate_checks_configuration_without_external_secrets(self) -> None: with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) stdout = StringIO() with contextlib.redirect_stdout(stdout): code = main(["validate", "--config", str(config)]) self.assertEqual(0, code) payload = json.loads(stdout.getvalue()) self.assertEqual("valid", payload["status"]) self.assertEqual(1, payload["targets"]) def test_watch_dry_run_needs_no_secrets_and_has_no_side_effects(self) -> None: with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) stdout = StringIO() with ( patch.dict("os.environ", {}, clear=True), patch("ai_triage.cli.GitHubClient") as github, patch("ai_triage.cli.StateStore") as state, contextlib.redirect_stdout(stdout), ): code = main(["watch", "--config", str(config), "--dry-run"]) self.assertEqual(0, code) payload = json.loads(stdout.getvalue()) self.assertTrue(payload["dry_run"]) self.assertEqual(["safe-example"], payload["targets"]) self.assertTrue(payload["state_path"].endswith("state/events.sqlite3")) github.assert_not_called() state.assert_not_called() def test_live_watch_with_no_enabled_targets_exits_without_secrets(self) -> None: with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory, enabled=False) stdout = StringIO() with patch.dict("os.environ", {}, clear=True), contextlib.redirect_stdout(stdout): code = main(["watch", "--config", str(config), "--once"]) self.assertEqual(0, code) self.assertEqual("no_targets", json.loads(stdout.getvalue())["status"]) def test_watch_once_returns_nonzero_and_prints_sanitized_cycle_errors(self) -> None: RecordingMonitor.result = CycleResult( errors=(CycleError("safe-example", None, "discover", "api unavailable"),) ) with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) stdout = StringIO() with self.live_patches(), contextlib.redirect_stdout(stdout): code = main(["watch", "--config", str(config), "--once"]) self.assertEqual(1, code) payload = json.loads(stdout.getvalue()) self.assertEqual("completed_with_errors", payload["status"]) self.assertEqual("discover", payload["errors"][0]["stage"]) def test_cli_overrides_state_and_output_paths(self) -> None: with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) state_path = Path(directory, "custom", "state.sqlite3") output_path = Path(directory, "custom", "reports") with self.live_patches(state_type=RecordingState): code = main( [ "watch", "--config", str(config), "--once", "--state-path", str(state_path), "--output-dir", str(output_path), ] ) self.assertEqual(0, code) monitor = RecordingMonitor.instances[-1] self.assertEqual(output_path, monitor.kwargs["output_root"]) self.assertEqual(state_path, monitor.args[1].path) def test_continuous_watch_passes_signal_stop_event_to_monitor(self) -> None: with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) with self.live_patches(state_type=RecordingState): code = main(["watch", "--config", str(config)]) self.assertEqual(0, code) monitor = RecordingMonitor.instances[-1] self.assertTrue(monitor.ran) self.assertTrue(monitor.stop_event.is_set()) self.assertIsNotNone(monitor.on_cycle) def test_missing_live_credentials_are_configuration_errors(self) -> None: with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) for environment, missing_name in ( ({}, "AI_TRIAGE_GITHUB_TOKEN"), ({"AI_TRIAGE_GITHUB_TOKEN": "token"}, "AI_TRIAGE_FEISHU_WEBHOOK"), ( { "AI_TRIAGE_GITHUB_TOKEN": "token", "AI_TRIAGE_FEISHU_WEBHOOK": "https://feishu.invalid/hook", }, "LARK_APP_ID", ), ( { "AI_TRIAGE_GITHUB_TOKEN": "token", "AI_TRIAGE_FEISHU_WEBHOOK": "https://feishu.invalid/hook", "LARK_APP_ID": "app-id", }, "LARK_APP_SECRET", ), ( { "AI_TRIAGE_GITHUB_TOKEN": "token", "AI_TRIAGE_FEISHU_WEBHOOK": "https://feishu.invalid/hook", "LARK_APP_ID": "app-id", "LARK_APP_SECRET": "app-secret", }, "ONCALL_API_KEY", ), ): with self.subTest(missing=missing_name): stderr = StringIO() with ( patch.dict("os.environ", environment, clear=True), contextlib.redirect_stderr(stderr), ): code = main(["watch", "--config", str(config), "--once"]) self.assertEqual(2, code) self.assertIn(missing_name, stderr.getvalue()) def test_live_watch_resolves_owner_mentions_for_notifier(self) -> None: with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) with ( self.live_patches(), patch( "ai_triage.cli.resolve_owner_mentions", return_value={"case": "ou_case"}, ) as resolver, patch("ai_triage.cli.FeishuWebhookNotifier") as notifier, ): code = main(["watch", "--config", str(config), "--once"]) self.assertEqual(0, code) resolver.assert_called_once() notifier.assert_called_once_with( "https://feishu.invalid/hook", owner_mentions={"case": "ou_case"}, ) def test_live_watch_uses_separate_acp_protocol_limit(self) -> None: with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) with ( self.live_patches(), patch("ai_triage.cli.ACPClient", return_value=object()) as acp, ): code = main(["watch", "--config", str(config), "--once"]) self.assertEqual(0, code) acp.assert_called_once_with( ("opencode", "acp"), timeout_seconds=60.0, max_output_bytes=500_000, ) def test_live_watch_loads_api_team_mapping_for_triage_engine(self) -> None: router = object() with tempfile.TemporaryDirectory() as directory: config = self.fixture(directory) with ( self.live_patches(), patch( "ai_triage.cli.load_api_team_router", return_value=router, ) as loader, patch("ai_triage.cli.TriageEngine", return_value=object()) as engine, ): code = main(["watch", "--config", str(config), "--once"]) self.assertEqual(0, code) mapping_path = ( Path(directory).resolve() / "routing" / "api-team-map.toml" ) self.assertEqual( [ call(mapping_path), call( mapping_path, {"global_control": "SOE - Global & Dedicated"}, ), ], loader.call_args_list, ) self.assertIs(engine.call_args.kwargs["api_team_router"], router) def live_patches(self, *, state_type: type = RecordingState): stack = contextlib.ExitStack() stack.enter_context( patch.dict( "os.environ", { "AI_TRIAGE_GITHUB_TOKEN": "synthetic-token", "AI_TRIAGE_FEISHU_WEBHOOK": "https://feishu.invalid/hook", "LARK_APP_ID": "synthetic-app-id", "LARK_APP_SECRET": "synthetic-app-secret", "ONCALL_API_KEY": "synthetic-oncall-key", }, clear=True, ) ) stack.enter_context( patch("ai_triage.cli.OncallClient", return_value=object()) ) stack.enter_context( patch( "ai_triage.cli.load_oncall_teams", return_value={"global_control": "SOE - Global & Dedicated"}, ) ) stack.enter_context( patch("ai_triage.cli.OncallOwnerRouter", return_value=object()) ) stack.enter_context( patch("ai_triage.cli.FeishuUserResolver") ) stack.enter_context(patch("ai_triage.cli.GitHubClient", return_value=object())) stack.enter_context(patch("ai_triage.cli.ACPClient", return_value=object())) stack.enter_context(patch("ai_triage.cli.TriageEngine", return_value=object())) stack.enter_context(patch("ai_triage.cli.Orchestrator", return_value=object())) stack.enter_context( patch("ai_triage.cli.FeishuWebhookNotifier", return_value=object()) ) stack.enter_context( patch( "ai_triage.cli.resolve_owner_mentions", return_value={"case": "ou_case"}, ) ) stack.enter_context( patch("ai_triage.cli.StateStore", side_effect=state_type) ) stack.enter_context(patch("ai_triage.cli.Monitor", RecordingMonitor)) return stack if __name__ == "__main__": unittest.main()