Fix unit tests

Author: tolik0

airbyte_cdk/sources/declarative/incremental/concurrent_partition_cursor.py

@@ -99,23 +99,29 @@ def __init__(
         self._semaphore_per_partition: OrderedDict[str, threading.Semaphore] = OrderedDict()
 
         # Parent-state tracking: store each partition’s parent state in creation order
-        self._partition_parent_state_map: OrderedDict[str, tuple[Mapping[str, Any], int]] = OrderedDict()
+        self._partition_parent_state_map: OrderedDict[str, tuple[Mapping[str, Any], int]] = (
+            OrderedDict()
+        )
+        self._parent_state: Optional[StreamState] = None
 
+        # Tracks when the last slice for partition is emitted
         self._finished_partitions: set[str] = set()
+        # Used to track the sequence numbers of open partitions
         self._open_seqs: deque[int] = deque()
         self._next_seq: int = 0
+        # Dictionary to map partition keys to their sequence numbers
         self._seq_by_partition: dict[str, int] = {}
 
         self._lock = threading.Lock()
-        self._timer = Timer()
-        self._new_global_cursor: Optional[StreamState] = None
         self._lookback_window: int = 0
-        self._parent_state: Optional[StreamState] = None
+        self._new_global_cursor: Optional[StreamState] = None
         self._number_of_partitions: int = 0
         self._use_global_cursor: bool = use_global_cursor
         self._partition_serializer = PerPartitionKeySerializer()
+
         # Track the last time a state message was emitted
         self._last_emission_time: float = 0.0
+        self._timer = Timer()
 
         self._set_initial_state(stream_state)
 
@@ -177,8 +183,9 @@ def _check_and_update_parent_state(self) -> None:
         last_closed_state = None
 
         while self._partition_parent_state_map:
-            earliest_key, (candidate_state, candidate_seq) = \
-                next(iter(self._partition_parent_state_map.items()))
+            earliest_key, (candidate_state, candidate_seq) = next(
+                iter(self._partition_parent_state_map.items())
+            )
 
             # if any partition that started <= candidate_seq is still open, we must wait
             if self._open_seqs and self._open_seqs[0] <= candidate_seq:
@@ -459,10 +466,10 @@ def _cleanup_if_done(self, partition_key: str) -> None:
         cursor, semaphore, flag inside `_finished_partitions`
         """
         if not (
-                partition_key in self._finished_partitions
-                and self._semaphore_per_partition[partition_key]._value == 0
+            partition_key in self._finished_partitions
+            and self._semaphore_per_partition[partition_key]._value == 0
         ):
-           return
+            return
 
         self._semaphore_per_partition.pop(partition_key, None)
         self._finished_partitions.discard(partition_key)

unit_tests/sources/declarative/incremental/test_concurrent_perpartitioncursor.py

@@ -3436,8 +3436,12 @@ def test_given_unfinished_first_parent_partition_no_parent_state_update():
     }
     assert mock_cursor_1.stream_slices.call_count == 1  # Called once for each partition
     assert mock_cursor_2.stream_slices.call_count == 1  # Called once for each partition
-    assert len(cursor._semaphore_per_partition) == 1  # Semaphore cleaned after partiton is completed
-    # ToDo: Add check for other interal values
+
+    assert len(cursor._semaphore_per_partition) == 1
+    assert len(cursor._finished_partitions) == 1
+    assert len(cursor._open_seqs) == 1
+    assert len(cursor._seq_by_partition) == 1
+
 
 def test_given_unfinished_last_parent_partition_with_partial_parent_state_update():
     # Create two mock cursors with different states for each partition
@@ -3520,8 +3524,11 @@ def test_given_unfinished_last_parent_partition_with_partial_parent_state_update
     }
     assert mock_cursor_1.stream_slices.call_count == 1  # Called once for each partition
     assert mock_cursor_2.stream_slices.call_count == 1  # Called once for each partition
-    assert len(cursor._semaphore_per_partition) == 0
-    # ToDo: Add check for other interal values
+
+    assert len(cursor._semaphore_per_partition) == 1
+    assert len(cursor._finished_partitions) == 1
+    assert len(cursor._open_seqs) == 1
+    assert len(cursor._seq_by_partition) == 1
 
 
 def test_given_all_partitions_finished_when_close_partition_then_final_state_emitted():
@@ -3596,7 +3603,12 @@ def test_given_all_partitions_finished_when_close_partition_then_final_state_emi
     assert final_state["lookback_window"] == 1
     assert cursor._message_repository.emit_message.call_count == 2
     assert mock_cursor.stream_slices.call_count == 2  # Called once for each partition
-    assert len(cursor._semaphore_per_partition) == 1
+
+    # Checks that all internal variables are cleaned up
+    assert len(cursor._semaphore_per_partition) == 0
+    assert len(cursor._finished_partitions) == 0
+    assert len(cursor._open_seqs) == 0
+    assert len(cursor._seq_by_partition) == 0
 
 
 def test_given_partition_limit_exceeded_when_close_partition_then_switch_to_global_cursor():
@@ -3715,18 +3727,20 @@ def test_semaphore_cleanup():
     # Verify initial state
     assert len(cursor._semaphore_per_partition) == 2
     assert len(cursor._partition_parent_state_map) == 2
-    assert cursor._partition_parent_state_map['{"id":"1"}'] == {"parent": {"state": "state1"}}
-    assert cursor._partition_parent_state_map['{"id":"2"}'] == {"parent": {"state": "state2"}}
+    assert len(cursor._open_seqs) == 2
+    assert len(cursor._seq_by_partition) == 2
+    assert cursor._partition_parent_state_map['{"id":"1"}'][0] == {"parent": {"state": "state1"}}
+    assert cursor._partition_parent_state_map['{"id":"2"}'][0] == {"parent": {"state": "state2"}}
 
     # Close partitions to acquire semaphores (value back to 0)
     for s in generated_slices:
         cursor.close_partition(DeclarativePartition("test_stream", {}, MagicMock(), MagicMock(), s))
 
     # Check state after closing partitions
-    assert len(cursor._finished_partitions) == 2
+    assert len(cursor._finished_partitions) == 0
     assert len(cursor._semaphore_per_partition) == 0
-    assert '{"id":"1"}' not in cursor._semaphore_per_partition
-    assert '{"id":"2"}' not in cursor._semaphore_per_partition
+    assert len(cursor._open_seqs) == 0
+    assert len(cursor._seq_by_partition) == 0
     assert len(cursor._partition_parent_state_map) == 0  # All parent states should be popped
     assert cursor._parent_state == {"parent": {"state": "state2"}}  # Last parent state