add timeout mcp poll
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
import unittest
|
||||
import time
|
||||
|
||||
from fgai.graylog_aggregate import GraylogAggregateSource
|
||||
|
||||
@@ -92,6 +93,17 @@ class GraylogAggregateTests(unittest.TestCase):
|
||||
self.assertIn("groupings", client.arguments[0])
|
||||
self.assertNotIn("group_by", client.arguments[0])
|
||||
|
||||
def test_deadline_stops_aggregate_attempts(self):
|
||||
client = _AggregateClient([
|
||||
{"result": {"content": [{"type": "text", "text": '{"events": 9}'}]}}
|
||||
])
|
||||
|
||||
status = GraylogAggregateSource(client, "firewall").fetch_count(deadline_monotonic=time.monotonic() - 1)
|
||||
|
||||
self.assertEqual(status["aggregate_status"], "error")
|
||||
self.assertIn("poll_budget_exceeded", status["aggregate_error"])
|
||||
self.assertEqual(client.arguments, [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import unittest
|
||||
import time
|
||||
|
||||
from fgai.graylog_source import GraylogStreamSource
|
||||
|
||||
@@ -73,6 +74,14 @@ class GraylogSourceTests(unittest.TestCase):
|
||||
self.assertTrue(status["partial"])
|
||||
self.assertIn("graylog_search_error", status["error"])
|
||||
|
||||
def test_deadline_stops_raw_fetch_before_search(self):
|
||||
client = _Client()
|
||||
events, status = GraylogStreamSource(client, "vpn").fetch(deadline_monotonic=time.monotonic() - 1)
|
||||
self.assertEqual(events, [])
|
||||
self.assertTrue(status["partial"])
|
||||
self.assertEqual(status["error"], "poll_budget_exceeded")
|
||||
self.assertIsNone(client.arguments)
|
||||
|
||||
def test_returns_partial_status_instead_of_raising_on_probe_error(self):
|
||||
events, status = GraylogStreamSource(_ProbeErrorClient(), "vpn").fetch()
|
||||
self.assertEqual(events, [])
|
||||
|
||||
Reference in New Issue
Block a user