fix mcp

tests/test_graylog_aggregate.py

@@ -20,6 +20,11 @@ class _AggregateClient:
         return self.responses.pop(0)
 
 
+class _ProbeErrorClient:
+    def probe(self):
+        raise RuntimeError("connection refused")
+
+
 class GraylogAggregateTests(unittest.TestCase):
     def test_reads_count_from_graylog_schema_rows(self):
         client = _AggregateClient([
@@ -33,6 +38,11 @@ class GraylogAggregateTests(unittest.TestCase):
         self.assertEqual(client.arguments[0]["streams"], ["firewall"])
         self.assertEqual(client.arguments[0]["metrics"], ["count()"])
 
+    def test_returns_error_status_instead_of_raising_on_probe_error(self):
+        status = GraylogAggregateSource(_ProbeErrorClient(), "firewall").fetch_count()
+        self.assertEqual(status["aggregate_status"], "error")
+        self.assertIn("probe_error", status["aggregate_error"])
+
     def test_tries_fallback_argument_shape_after_tool_error(self):
         client = _AggregateClient([
             {"result": {"isError": True, "content": [{"type": "text", "text": "bad metrics"}]}},

tests/test_graylog_source.py

@@ -23,6 +23,11 @@ class _ErrorClient:
         return {"result": {"isError": True, "content": [{"type": "text", "text": "Tool call failed: timeout"}]}}
 
 
+class _ProbeErrorClient:
+    def probe(self):
+        raise RuntimeError("connection refused")
+
+
 class GraylogSourceTests(unittest.TestCase):
     def test_applies_custom_mapping_to_generic_stream_message(self):
         client = _Client()
@@ -56,6 +61,12 @@ class GraylogSourceTests(unittest.TestCase):
         self.assertTrue(status["partial"])
         self.assertIn("graylog_search_error", status["error"])
 
+    def test_returns_partial_status_instead_of_raising_on_probe_error(self):
+        events, status = GraylogStreamSource(_ProbeErrorClient(), "vpn").fetch()
+        self.assertEqual(events, [])
+        self.assertTrue(status["partial"])
+        self.assertIn("probe_error", status["error"])
+
     def test_requests_selected_profile_fields(self):
         client = _Client()
         GraylogStreamSource(client, "windows", profile_fields=("TargetUserName", "EventID")).fetch()