Fix unit tests

Author: tolik0

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -171,10 +171,21 @@ def ensure_at_least_one_state_emitted(self) -> None:
             self._parent_state = self._partition_router.get_stream_state()
         self._emit_state_message(throttle=False)
 
-    def _emit_state_message(self, throttle: bool = True) -> None:
+    def _throttle_state_message(self) -> Optional[float]:
+        """
+        Throttles the state message emission to once every 60 seconds.
+        """
         current_time = time.time()
-        if throttle and current_time - self._last_emission_time <= 60:
-            return
+        if current_time - self._last_emission_time <= 60:
+            return None
+        return current_time
+
+    def _emit_state_message(self, throttle: bool = True) -> None:
+        if throttle:
+            current_time = self._throttle_state_message()
+            if current_time is None:
+                return
+            self._last_emission_time = current_time
         self._connector_state_manager.update_state_for_stream(
             self._stream_name,
             self._stream_namespace,
@@ -184,7 +195,6 @@ def _emit_state_message(self, throttle: bool = True) -> None:
             self._stream_name, self._stream_namespace
         )
         self._message_repository.emit_message(state_message)
-        self._last_emission_time = current_time
 
     def stream_slices(self) -> Iterable[StreamSlice]:
         if self._timer.is_running():
@@ -358,6 +368,7 @@ def _set_global_state(self, stream_state: Mapping[str, Any]) -> None:
             self._new_global_cursor = deepcopy(fixed_global_state)
 
     def observe(self, record: Record) -> None:
+        # ToDo: check number of partitions
         if not self._use_global_cursor and self.limit_reached():
             logger.info(
                 f"Exceeded the 'SWITCH_TO_GLOBAL_LIMIT' of {self.SWITCH_TO_GLOBAL_LIMIT}. "

unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py

@@ -3,6 +3,7 @@
 from copy import deepcopy
 from datetime import datetime, timedelta
 from typing import Any, List, Mapping, MutableMapping, Optional, Union
+from unittest.mock import MagicMock, patch
 from urllib.parse import unquote
 
 import pytest
@@ -18,6 +19,7 @@
 from airbyte_cdk.sources.declarative.concurrent_declarative_source import (
     ConcurrentDeclarativeSource,
 )
+from airbyte_cdk.sources.declarative.incremental import ConcurrentPerPartitionCursor
 from airbyte_cdk.test.catalog_builder import CatalogBuilder, ConfiguredAirbyteStreamBuilder
 from airbyte_cdk.test.entrypoint_wrapper import EntrypointOutput, read
 
@@ -1181,14 +1183,18 @@ def test_incremental_parent_state(
     initial_state,
     expected_state,
 ):
-    run_incremental_parent_state_test(
-        manifest,
-        mock_requests,
-        expected_records,
-        num_intermediate_states,
-        initial_state,
-        [expected_state],
-    )
+    # Patch `_throttle_state_message` so it always returns a float (indicating "no throttle")
+    with patch.object(
+        ConcurrentPerPartitionCursor, "_throttle_state_message", return_value=9999999.0
+    ):
+        run_incremental_parent_state_test(
+            manifest,
+            mock_requests,
+            expected_records,
+            num_intermediate_states,
+            initial_state,
+            [expected_state],
+        )
 
 
 STATE_MIGRATION_EXPECTED_STATE = {
@@ -2967,3 +2973,47 @@ def test_incremental_substream_request_options_provider(
         expected_records,
         expected_state,
     )
+
+
+def test_state_throttling(mocker):
+    """
+    Verifies that _emit_state_message does not emit a new state if less than 60s
+    have passed since last emission, and does emit once 60s or more have passed.
+    """
+    cursor = ConcurrentPerPartitionCursor(
+        cursor_factory=MagicMock(),
+        partition_router=MagicMock(),
+        stream_name="test_stream",
+        stream_namespace=None,
+        stream_state={},
+        message_repository=MagicMock(),
+        connector_state_manager=MagicMock(),
+        connector_state_converter=MagicMock(),
+        cursor_field=MagicMock(),
+    )
+
+    mock_connector_manager = cursor._connector_state_manager
+    mock_repo = cursor._message_repository
+
+    # Set the last emission time to "0" so we can control offset from that
+    cursor._last_emission_time = 0
+
+    mock_time = mocker.patch("time.time")
+
+    # First attempt: only 10 seconds passed => NO emission
+    mock_time.return_value = 10
+    cursor._emit_state_message()
+    mock_connector_manager.update_state_for_stream.assert_not_called()
+    mock_repo.emit_message.assert_not_called()
+
+    # Second attempt: 30 seconds passed => still NO emission
+    mock_time.return_value = 30
+    cursor._emit_state_message()
+    mock_connector_manager.update_state_for_stream.assert_not_called()
+    mock_repo.emit_message.assert_not_called()
+
+    # Advance time: 70 seconds => exceed 60s => MUST emit
+    mock_time.return_value = 70
+    cursor._emit_state_message()
+    mock_connector_manager.update_state_for_stream.assert_called_once()
+    mock_repo.emit_message.assert_called_once()