Ubika : Add timestepper to connector

Ubika/CHANGELOG.md

@@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+## 2025-12-09 - 1.0.8
+
+### Added
+
+- Add timestepper to avoid timeouts on large data fetches
+- 
 ## 2026-02-03 - 1.0.7
 
 ### Fixed

Ubika/manifest.json

@@ -11,7 +11,7 @@
   "name": "Ubika",
   "uuid": "0c82ee9b-f645-47f9-8e16-a689cfc246c4",
   "slug": "ubika",
-  "version": "1.0.7",
+  "version": "1.0.8",
   "categories": [
     "Network"
   ]

Ubika/tests/test_timestepper.py

@@ -0,0 +1,244 @@
+import datetime
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ubika_modules.timestepper import TimeStepper
+
+
+@pytest.fixture
+def mock_trigger():
+    trigger = MagicMock()
+    trigger.log = MagicMock()
+    trigger.configuration.intake_key = "test_intake_key"
+    return trigger
+
+
+@pytest.fixture
+def fixed_time():
+    return datetime.datetime(2024, 1, 15, 12, 0, 0, tzinfo=datetime.UTC)
+
+
+def test_timestepper_initialization(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 10, 0, 0, tzinfo=datetime.UTC)
+    end = datetime.datetime(2024, 1, 15, 11, 0, 0, tzinfo=datetime.UTC)
+    frequency = datetime.timedelta(seconds=60)
+    timedelta = datetime.timedelta(minutes=1)
+
+    stepper = TimeStepper(mock_trigger, start, end, frequency, timedelta)
+
+    assert stepper.trigger == mock_trigger
+    assert stepper.start == start
+    assert stepper.end == end
+    assert stepper.frequency == frequency
+    assert stepper.timedelta == timedelta
+
+
+def test_ranges_yields_current_time_range_first(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 10, 0, 0, tzinfo=datetime.UTC)
+    end = datetime.datetime(2024, 1, 15, 11, 0, 0, tzinfo=datetime.UTC)
+    frequency = datetime.timedelta(seconds=60)
+    timedelta = datetime.timedelta(minutes=1)
+
+    stepper = TimeStepper(mock_trigger, start, end, frequency, timedelta)
+
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime, patch("ubika_modules.timestepper.time"):
+        mock_datetime.datetime.now.return_value = datetime.datetime(2024, 1, 15, 12, 0, 0, tzinfo=datetime.UTC)
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+
+        generator = stepper.ranges()
+        first_range = next(generator)
+
+        assert first_range == (start, end)
+
+
+def test_ranges_updates_start_and_end_after_iteration(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 10, 0, 0, tzinfo=datetime.UTC)
+    end = datetime.datetime(2024, 1, 15, 10, 1, 0, tzinfo=datetime.UTC)
+    frequency = datetime.timedelta(seconds=60)
+    timedelta = datetime.timedelta(minutes=1)
+
+    stepper = TimeStepper(mock_trigger, start, end, frequency, timedelta)
+
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime, patch("ubika_modules.timestepper.time"):
+        mock_datetime.datetime.now.return_value = datetime.datetime(2024, 1, 15, 12, 0, 0, tzinfo=datetime.UTC)
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+
+        generator = stepper.ranges()
+        next(generator)
+        second_range = next(generator)
+
+        assert second_range[0] == end
+        assert second_range[1] > end
+
+
+def test_ranges_sleeps_when_next_end_is_in_future(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 11, 58, 0, tzinfo=datetime.UTC)
+    end = datetime.datetime(2024, 1, 15, 11, 59, 0, tzinfo=datetime.UTC)
+    frequency = datetime.timedelta(seconds=60)
+    timedelta = datetime.timedelta(minutes=1)
+
+    stepper = TimeStepper(mock_trigger, start, end, frequency, timedelta)
+
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime, patch(
+        "ubika_modules.timestepper.time"
+    ) as mock_time:
+        now = datetime.datetime(2024, 1, 15, 11, 59, 30, tzinfo=datetime.UTC)
+        mock_datetime.datetime.now.return_value = now
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+        mock_time.sleep = MagicMock()
+
+        generator = stepper.ranges()
+        next(generator)
+        next(generator)
+
+        assert mock_time.sleep.call_count == 1
+        sleep_duration = mock_time.sleep.call_args[0][0]
+        assert sleep_duration > 0
+
+
+def test_ranges_does_not_sleep_when_next_end_is_in_past(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 10, 0, 0, tzinfo=datetime.UTC)
+    end = datetime.datetime(2024, 1, 15, 10, 1, 0, tzinfo=datetime.UTC)
+    frequency = datetime.timedelta(seconds=60)
+    timedelta = datetime.timedelta(minutes=1)
+
+    stepper = TimeStepper(mock_trigger, start, end, frequency, timedelta)
+
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime, patch(
+        "ubika_modules.timestepper.time"
+    ) as mock_time:
+        mock_datetime.datetime.now.return_value = datetime.datetime(2024, 1, 15, 12, 0, 0, tzinfo=datetime.UTC)
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+        mock_time.sleep = MagicMock()
+
+        generator = stepper.ranges()
+        next(generator)
+        next(generator)
+
+        assert mock_time.sleep.call_count == 0
+
+
+def test_create_with_default_parameters(mock_trigger, fixed_time):
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime:
+        mock_datetime.datetime.now.return_value = fixed_time
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+
+        stepper = TimeStepper.create(mock_trigger)
+
+        assert stepper.trigger == mock_trigger
+        assert stepper.frequency == datetime.timedelta(seconds=60)
+        assert stepper.timedelta == datetime.timedelta(minutes=1)
+        assert stepper.end < fixed_time
+        assert stepper.start < stepper.end
+
+
+def test_create_with_start_time_zero(mock_trigger, fixed_time):
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime:
+        mock_datetime.datetime.now.return_value = fixed_time
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+
+        stepper = TimeStepper.create(mock_trigger, start_time=0)
+
+        expected_end = fixed_time - datetime.timedelta(minutes=1)
+        assert stepper.end == expected_end
+
+
+def test_create_with_custom_start_time(mock_trigger, fixed_time):
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime:
+        mock_datetime.datetime.now.return_value = fixed_time
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+
+        stepper = TimeStepper.create(mock_trigger, start_time=2)
+
+        expected_end = fixed_time - datetime.timedelta(hours=2)
+        assert stepper.end == expected_end
+
+
+def test_create_with_custom_frequency_and_timedelta(mock_trigger, fixed_time):
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime:
+        mock_datetime.datetime.now.return_value = fixed_time
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+
+        stepper = TimeStepper.create(mock_trigger, frequency=120, timedelta=5)
+
+        assert stepper.frequency == datetime.timedelta(seconds=120)
+        assert stepper.timedelta == datetime.timedelta(minutes=5)
+
+
+def test_create_from_time_with_default_parameters(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 10, 0, 0, tzinfo=datetime.UTC)
+
+    stepper = TimeStepper.create_from_time(mock_trigger, start)
+
+    assert stepper.trigger == mock_trigger
+    assert stepper.start == start
+    assert stepper.end == start + datetime.timedelta(seconds=60)
+    assert stepper.frequency == datetime.timedelta(seconds=60)
+    assert stepper.timedelta == datetime.timedelta(minutes=1)
+
+
+def test_create_from_time_with_custom_parameters(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 10, 0, 0, tzinfo=datetime.UTC)
+
+    stepper = TimeStepper.create_from_time(mock_trigger, start, frequency=300, timedelta=10)
+
+    assert stepper.start == start
+    assert stepper.end == start + datetime.timedelta(seconds=300)
+    assert stepper.frequency == datetime.timedelta(seconds=300)
+    assert stepper.timedelta == datetime.timedelta(minutes=10)
+
+
+def test_ranges_logs_current_lag(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 10, 0, 0, tzinfo=datetime.UTC)
+    end = datetime.datetime(2024, 1, 15, 10, 1, 0, tzinfo=datetime.UTC)
+    frequency = datetime.timedelta(seconds=60)
+    timedelta = datetime.timedelta(minutes=1)
+
+    stepper = TimeStepper(mock_trigger, start, end, frequency, timedelta)
+
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime, patch("ubika_modules.timestepper.time"):
+        mock_datetime.datetime.now.return_value = datetime.datetime(2024, 1, 15, 12, 0, 0, tzinfo=datetime.UTC)
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+
+        generator = stepper.ranges()
+        next(generator)
+        next(generator)
+
+        assert mock_trigger.log.called
+        log_calls = [call for call in mock_trigger.log.call_args_list if "Current lag" in str(call)]
+        assert len(log_calls) > 0
+
+
+def test_ranges_logs_waiting_message_when_sleeping(mock_trigger):
+    start = datetime.datetime(2024, 1, 15, 11, 58, 0, tzinfo=datetime.UTC)
+    end = datetime.datetime(2024, 1, 15, 11, 59, 0, tzinfo=datetime.UTC)
+    frequency = datetime.timedelta(seconds=60)
+    timedelta = datetime.timedelta(minutes=1)
+
+    stepper = TimeStepper(mock_trigger, start, end, frequency, timedelta)
+
+    with patch("ubika_modules.timestepper.datetime") as mock_datetime, patch(
+        "ubika_modules.timestepper.time"
+    ) as mock_time:
+        now = datetime.datetime(2024, 1, 15, 11, 59, 30, tzinfo=datetime.UTC)
+        mock_datetime.datetime.now.return_value = now
+        mock_datetime.UTC = datetime.UTC
+        mock_datetime.timedelta = datetime.timedelta
+        mock_time.sleep = MagicMock()
+
+        generator = stepper.ranges()
+        next(generator)
+        next(generator)
+
+        log_calls = [call for call in mock_trigger.log.call_args_list if "Waiting" in str(call)]
+        assert len(log_calls) > 0

Ubika/tests/test_ubika_cloud_protector_next_gen.py

@@ -1,4 +1,5 @@
 from unittest.mock import MagicMock, patch
+from datetime import datetime, timezone
 
 import httpx
 import pytest
@@ -157,7 +158,7 @@ def test_fetch_events_with_pagination(respx_mock: MockRouter, trigger, message1,
         ).mock(return_value=httpx.Response(200, json=message2))
 
         trigger.from_timestamp = 1747326567845
-        events = trigger.fetch_events()
+        events = trigger._UbikaCloudProtectorNextGenConnector__fetch_next_events(1747326567845, 1747326667845)
 
         assert list(events) == [message1["spec"]["items"]]
 
@@ -178,10 +179,17 @@ def test_next_batch_sleep_until_next_round(respx_mock: MockRouter, trigger, mess
             )
         )
 
+        start = datetime.fromtimestamp(1747326560, tz=timezone.utc)
+        end = datetime.fromtimestamp(1747326660, tz=timezone.utc)
+
+        start_ms = int(start.timestamp() * 1000)
+        end_ms = int(end.timestamp() * 1000)
+
         respx_mock.get(
             "https://api.ubika.io/rest/logs.ubika.io/v1/ns/sekoia/security-events",
             params={
-                "filters.fromDate": "1747326567845",
+                "filters.fromDate": str(start_ms),
+                "filters.toDate": str(end_ms),
                 "pagination.realtime": "true",
                 "pagination.pageSize": "100",
             },
@@ -196,14 +204,13 @@ def test_next_batch_sleep_until_next_round(respx_mock: MockRouter, trigger, mess
             },
         ).mock(return_value=httpx.Response(200, json=message2))
 
-        trigger.from_timestamp = 1747326567845
-
         batch_duration = trigger.configuration.frequency + 20
         start_time = 1747326560
         end_time = start_time + batch_duration
         mock_time.time.side_effect = [start_time, end_time, end_time]
 
-        trigger.next_batch()
+
+        trigger.next_batch(start, end)
 
         assert trigger.push_events_to_intakes.call_count == 1
         assert mock_time.sleep.call_count == 0
@@ -219,8 +226,11 @@ def test_authorization_http_error_without_retry(respx_mock: MockRouter, trigger)
         )
     )
 
+    start = datetime.fromtimestamp(1747326560, tz=timezone.utc)
+    end = datetime.fromtimestamp(1747326660, tz=timezone.utc)
+
     with patch("time.sleep"), pytest.raises(AuthorizationError):
-        trigger.next_batch()
+        trigger.next_batch(start=start, end=end)
 
     assert route.call_count == 1
 
@@ -235,8 +245,11 @@ def test_authorization_http_error_with_retry(respx_mock: MockRouter, trigger):
         )
     )
 
+    start = datetime.fromtimestamp(1747326560, tz=timezone.utc)
+    end = datetime.fromtimestamp(1747326660, tz=timezone.utc)
+
     with patch("time.sleep"), pytest.raises(AuthorizationError):
-        trigger.next_batch()
+        trigger.next_batch(start=start, end=end)
 
     assert route.call_count == 5
 
@@ -246,7 +259,10 @@ def test_authorization_timeout_error(respx_mock: MockRouter, trigger):
     route = respx_mock.post("/auth/realms/main/protocol/openid-connect/token")
     route.side_effect = httpx.TimeoutException("Timeout")
 
+    start = datetime.fromtimestamp(1747326560, tz=timezone.utc)
+    end = datetime.fromtimestamp(1747326660, tz=timezone.utc)
+
     with patch("time.sleep"), pytest.raises(AuthorizationTimeoutError):
-        trigger.next_batch()
+        trigger.next_batch(start=start, end=end)
 
     assert route.call_count == 5