fix test to use PartitionReader where cursor close is now invoked

Author: brianjlai

unit_tests/sources/streams/test_stream_read.py

@@ -4,6 +4,7 @@
 
 import logging
 from copy import deepcopy
+from queue import Queue
 from typing import Any, Dict, Iterable, List, Mapping, MutableMapping, Optional, Union
 from unittest.mock import Mock
 
@@ -589,7 +590,10 @@ def test_concurrent_incremental_read_two_slices():
         *records_partition_2,
     ]
 
-    expected_state = _create_state_message(
+    expected_state_1 = _create_state_message(
+        "__mock_stream", {"1": {"created_at": slice_timestamp_1}}
+    )
+    expected_state_2 = _create_state_message(
         "__mock_stream",
         {"1": {"created_at": slice_timestamp_1}, "2": {"created_at": slice_timestamp_2}},
     )
@@ -617,26 +621,27 @@ def test_concurrent_incremental_read_two_slices():
     for record in expected_records:
         assert record in actual_records
 
-    # We need run on_record to update cursor with record cursor value
-    for record in actual_records:
-        list(
-            handler.on_record(
-                Record(
-                    data=record,
-                    stream_name="__mock_stream",
-                )
-            )
-        )
+    # We need to process partitions generated by a PartitionReader in order to trigger
+    # the ConcurrentCursor.close_partition() flow and validate state is updated with
+    # the observed record values
+    partition_reader = PartitionReader(queue=Mock(spec=Queue))
+    assert isinstance(stream, StreamFacade)
+    abstract_stream = stream._abstract_stream
+    for partition in abstract_stream.generate_partitions():
+        partition_reader.process_partition(partition=partition, cursor=cursor)
 
     assert len(actual_records) == len(expected_records)
 
-    # We don't have a real source that reads from the message_repository for state, so we read from the queue directly to verify
-    # the cursor observed records correctly and updated partition states
-    mock_partition = Mock()
-    cursor.close_partition(mock_partition)
     actual_state = [state for state in message_repository.consume_queue()]
-    assert len(actual_state) == 1
-    assert actual_state[0] == expected_state
+    assert len(actual_state) == 2
+    assert (
+        actual_state[0].state.stream.stream_state.__dict__
+        == expected_state_1.state.stream.stream_state.__dict__
+    )
+    assert (
+        actual_state[1].state.stream.stream_state.__dict__
+        == expected_state_2.state.stream.stream_state.__dict__
+    )
 
 
 def setup_stream_dependencies(configured_json_schema):