Fixed the MCP error handling for large pulls.
This commit is contained in:
@@ -15,6 +15,14 @@ class _Client:
|
||||
return {"result": {"content": [{"type": "text", "text": '{"schema":[{"field":"client"},{"field":"server"},{"field":"result"}],"datarows":[["10.0.0.1","8.8.8.8","deny"],["10.0.0.2","8.8.8.8","accept"]]}' }]}}
|
||||
|
||||
|
||||
class _ErrorClient:
|
||||
def probe(self):
|
||||
return {"status": "connected"}
|
||||
|
||||
def call_tool(self, _name, _arguments):
|
||||
return {"result": {"isError": True, "content": [{"type": "text", "text": "Tool call failed: timeout"}]}}
|
||||
|
||||
|
||||
class GraylogSourceTests(unittest.TestCase):
|
||||
def test_applies_custom_mapping_to_generic_stream_message(self):
|
||||
client = _Client()
|
||||
@@ -42,6 +50,12 @@ class GraylogSourceTests(unittest.TestCase):
|
||||
self.assertEqual(client.arguments["limit"], 2)
|
||||
self.assertTrue(status["truncated"])
|
||||
|
||||
def test_returns_partial_status_instead_of_raising_on_search_error(self):
|
||||
events, status = GraylogStreamSource(_ErrorClient(), "vpn").fetch()
|
||||
self.assertEqual(events, [])
|
||||
self.assertTrue(status["partial"])
|
||||
self.assertIn("graylog_search_error", status["error"])
|
||||
|
||||
def test_requests_selected_profile_fields(self):
|
||||
client = _Client()
|
||||
GraylogStreamSource(client, "windows", profile_fields=("TargetUserName", "EventID")).fetch()
|
||||
|
||||
Reference in New Issue
Block a user