airbyte-cdk 0.52.7__py3-none-any.whl → 0.52.8__py3-none-any.whl

unit_tests/sources/streams/concurrent/scenarios/test_concurrent_scenarios.py

@@ -11,6 +11,9 @@ from pytest import LogCaptureFixture
 from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenario
 from unit_tests.sources.file_based.test_scenarios import verify_discover, verify_read
 from unit_tests.sources.streams.concurrent.scenarios.stream_facade_scenarios import (
+    test_incremental_stream_with_many_slices_but_without_slice_boundaries,
+    test_incremental_stream_with_slice_boundaries,
+    test_incremental_stream_without_slice_boundaries,
     test_stream_facade_multiple_streams,
     test_stream_facade_raises_exception,
     test_stream_facade_single_stream,
@@ -43,6 +46,9 @@ scenarios = [
     test_stream_facade_single_stream_with_multiple_slices,
     test_stream_facade_single_stream_with_multiple_slices_with_concurrency_level_two,
     test_stream_facade_raises_exception,
+    test_incremental_stream_with_slice_boundaries,
+    test_incremental_stream_without_slice_boundaries,
+    test_incremental_stream_with_many_slices_but_without_slice_boundaries,
 ]
 
 

unit_tests/sources/streams/concurrent/scenarios/thread_based_concurrent_stream_source_builder.py

@@ -12,15 +12,31 @@ from airbyte_cdk.sources.message import MessageRepository
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import AbstractAvailabilityStrategy, StreamAvailability, StreamAvailable
+from airbyte_cdk.sources.streams.concurrent.cursor import NoopCursor
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.partition_generator import PartitionGenerator
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 from airbyte_cdk.sources.streams.concurrent.thread_based_concurrent_stream import ThreadBasedConcurrentStream
+from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.utils.slice_logger import SliceLogger
 from airbyte_protocol.models import ConfiguredAirbyteStream
 from unit_tests.sources.file_based.scenarios.scenario_builder import SourceBuilder
 
 
+class LegacyStream(Stream):
+    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
+        return None
+
+    def read_records(
+        self,
+        sync_mode: SyncMode,
+        cursor_field: Optional[List[str]] = None,
+        stream_slice: Optional[Mapping[str, Any]] = None,
+        stream_state: Optional[Mapping[str, Any]] = None,
+    ) -> Iterable[StreamData]:
+        yield from []
+
+
 class ConcurrentCdkSource(AbstractSource):
     def __init__(self, streams: List[ThreadBasedConcurrentStream], message_repository: Optional[MessageRepository]):
         self._streams = streams
@@ -31,7 +47,7 @@ class ConcurrentCdkSource(AbstractSource):
         return True, None
 
     def streams(self, config: Mapping[str, Any]) -> List[Stream]:
-        return [StreamFacade(s) for s in self._streams]
+        return [StreamFacade(s, LegacyStream(), NoopCursor()) for s in self._streams]
 
     def spec(self, *args: Any, **kwargs: Any) -> ConnectorSpecification:
         return ConnectorSpecification(connectionSpecification={})
@@ -40,7 +56,7 @@ class ConcurrentCdkSource(AbstractSource):
         return ConfiguredAirbyteCatalog(
             streams=[
                 ConfiguredAirbyteStream(
-                    stream=StreamFacade(s).as_airbyte_stream(),
+                    stream=StreamFacade(s, LegacyStream(), NoopCursor()).as_airbyte_stream(),
                     sync_mode=SyncMode.full_refresh,
                     destination_sync_mode=DestinationSyncMode.overwrite,
                 )
@@ -57,7 +73,7 @@ class InMemoryPartitionGenerator(PartitionGenerator):
     def __init__(self, partitions: List[Partition]):
         self._partitions = partitions
 
-    def generate(self, sync_mode: SyncMode) -> Iterable[Partition]:
+    def generate(self) -> Iterable[Partition]:
         yield from self._partitions
 
 

unit_tests/sources/streams/concurrent/test_adapters.py

@@ -17,10 +17,16 @@ from airbyte_cdk.sources.streams.concurrent.adapters import (
     StreamPartitionGenerator,
 )
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import STREAM_AVAILABLE, StreamAvailable, StreamUnavailable
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor, NoopCursor
 from airbyte_cdk.sources.streams.concurrent.exceptions import ExceptionWithDisplayMessage
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+from airbyte_cdk.sources.streams.core import Stream
 from airbyte_cdk.sources.utils.transform import TransformConfig, TypeTransformer
 
+_ANY_SYNC_MODE = SyncMode.full_refresh
+_ANY_STATE = {"state_key": "state_value"}
+_ANY_CURSOR_FIELD = ["a", "cursor", "key"]
+
 
 @pytest.mark.parametrize(
     "stream_availability, expected_available, expected_message",
@@ -71,11 +77,12 @@ def test_stream_partition_generator(sync_mode):
     stream_slices = [{"slice": 1}, {"slice": 2}]
     stream.stream_slices.return_value = stream_slices
 
-    partition_generator = StreamPartitionGenerator(stream, message_repository)
+    partition_generator = StreamPartitionGenerator(stream, message_repository, _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE)
 
-    partitions = list(partition_generator.generate(sync_mode))
+    partitions = list(partition_generator.generate())
     slices = [partition.to_slice() for partition in partitions]
     assert slices == stream_slices
+    stream.stream_slices.assert_called_once_with(sync_mode=_ANY_SYNC_MODE, cursor_field=_ANY_CURSOR_FIELD, stream_state=_ANY_STATE)
 
 
 @pytest.mark.parametrize(
@@ -95,7 +102,10 @@ def test_stream_partition(transformer, expected_records):
     stream.transformer = transformer
     message_repository = InMemoryMessageRepository()
     _slice = None
-    partition = StreamPartition(stream, _slice, message_repository)
+    sync_mode = SyncMode.full_refresh
+    cursor_field = None
+    state = None
+    partition = StreamPartition(stream, _slice, message_repository, sync_mode, cursor_field, state)
 
     a_log_message = AirbyteMessage(
         type=MessageType.LOG,
@@ -128,7 +138,8 @@ def test_stream_partition_raising_exception(exception_type, expected_display_mes
 
     message_repository = InMemoryMessageRepository()
     _slice = None
-    partition = StreamPartition(stream, _slice, message_repository)
+
+    partition = StreamPartition(stream, _slice, message_repository, _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE)
 
     stream.read_records.side_effect = Exception()
 
@@ -148,7 +159,7 @@ def test_stream_partition_raising_exception(exception_type, expected_display_mes
 def test_stream_partition_hash(_slice, expected_hash):
     stream = Mock()
     stream.name = "stream"
-    partition = StreamPartition(stream, _slice, Mock())
+    partition = StreamPartition(stream, _slice, Mock(), _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE)
 
     _hash = partition.__hash__()
     assert _hash == expected_hash
@@ -163,7 +174,9 @@ class StreamFacadeTest(unittest.TestCase):
             json_schema={"type": "object"},
             supported_sync_modes=[SyncMode.full_refresh],
         )
-        self._facade = StreamFacade(self._abstract_stream)
+        self._legacy_stream = Mock(spec=Stream)
+        self._cursor = Mock(spec=Cursor)
+        self._facade = StreamFacade(self._abstract_stream, self._legacy_stream, self._cursor)
         self._logger = Mock()
         self._source = Mock()
         self._max_workers = 10
@@ -191,8 +204,14 @@ class StreamFacadeTest(unittest.TestCase):
         assert self._facade.get_json_schema() == json_schema
         self._abstract_stream.get_json_schema.assert_called_once_with()
 
-    def test_supports_incremental_is_false(self):
-        assert self._facade.supports_incremental is False
+    def test_given_cursor_is_noop_when_supports_incremental_then_return_legacy_stream_response(self):
+        assert (
+            StreamFacade(self._abstract_stream, self._legacy_stream, Mock(spec=NoopCursor)).supports_incremental
+            == self._legacy_stream.supports_incremental
+        )
+
+    def test_given_cursor_is_not_noop_when_supports_incremental_then_return_true(self):
+        assert StreamFacade(self._abstract_stream, self._legacy_stream, Mock(spec=Cursor)).supports_incremental
 
     def test_check_availability_is_delegated_to_wrapped_stream(self):
         availability = StreamAvailable()
@@ -219,8 +238,13 @@ class StreamFacadeTest(unittest.TestCase):
         assert actual_stream_data == expected_stream_data
 
     def test_read_records_incremental(self):
-        with self.assertRaises(NotImplementedError):
-            list(self._facade.read_records(SyncMode.incremental, None, None, None))
+        expected_stream_data = [{"data": 1}, {"data": 2}]
+        records = [Record(data) for data in expected_stream_data]
+        self._abstract_stream.read.return_value = records
+
+        actual_stream_data = list(self._facade.read_incremental(None, None, None, None, None, None, None))
+
+        assert actual_stream_data == expected_stream_data
 
     def test_create_from_stream_stream(self):
         stream = Mock()
@@ -228,7 +252,7 @@ class StreamFacadeTest(unittest.TestCase):
         stream.primary_key = "id"
         stream.cursor_field = "cursor"
 
-        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers)
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
 
         assert facade.name == "stream"
         assert facade.cursor_field == "cursor"
@@ -240,7 +264,7 @@ class StreamFacadeTest(unittest.TestCase):
         stream.primary_key = None
         stream.cursor_field = []
 
-        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers)
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
         facade._abstract_stream._primary_key is None
 
     def test_create_from_stream_with_composite_primary_key(self):
@@ -249,7 +273,7 @@ class StreamFacadeTest(unittest.TestCase):
         stream.primary_key = ["id", "name"]
         stream.cursor_field = []
 
-        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers)
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
         facade._abstract_stream._primary_key == ["id", "name"]
 
     def test_create_from_stream_with_empty_list_cursor(self):
@@ -257,7 +281,7 @@ class StreamFacadeTest(unittest.TestCase):
         stream.primary_key = "id"
         stream.cursor_field = []
 
-        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers)
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
 
         assert facade.cursor_field == []
 
@@ -267,7 +291,7 @@ class StreamFacadeTest(unittest.TestCase):
         stream.primary_key = [["field", "id"]]
 
         with self.assertRaises(ValueError):
-            StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers)
+            StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
 
     def test_create_from_stream_raises_exception_if_primary_key_has_invalid_type(self):
         stream = Mock()
@@ -275,7 +299,7 @@ class StreamFacadeTest(unittest.TestCase):
         stream.primary_key = 123
 
         with self.assertRaises(ValueError):
-            StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers)
+            StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
 
     def test_create_from_stream_raises_exception_if_cursor_field_is_nested(self):
         stream = Mock()
@@ -284,7 +308,7 @@ class StreamFacadeTest(unittest.TestCase):
         stream.cursor_field = ["field", "cursor"]
 
         with self.assertRaises(ValueError):
-            StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers)
+            StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
 
     def test_create_from_stream_with_cursor_field_as_list(self):
         stream = Mock()
@@ -292,7 +316,7 @@ class StreamFacadeTest(unittest.TestCase):
         stream.primary_key = "id"
         stream.cursor_field = ["cursor"]
 
-        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers)
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
         assert facade.cursor_field == "cursor"
 
     def test_create_from_stream_none_message_repository(self):
@@ -302,12 +326,12 @@ class StreamFacadeTest(unittest.TestCase):
         self._source.message_repository = None
 
         with self.assertRaises(ValueError):
-            StreamFacade.create_from_stream(self._stream, self._source, self._logger, self._max_workers)
+            StreamFacade.create_from_stream(self._stream, self._source, self._logger, self._max_workers, {}, self._cursor)
 
     def test_get_error_display_message_no_display_message(self):
         self._stream.get_error_display_message.return_value = "display_message"
 
-        facade = StreamFacade.create_from_stream(self._stream, self._source, self._logger, self._max_workers)
+        facade = StreamFacade.create_from_stream(self._stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
 
         expected_display_message = None
         e = Exception()
@@ -319,7 +343,7 @@ class StreamFacadeTest(unittest.TestCase):
     def test_get_error_display_message_with_display_message(self):
         self._stream.get_error_display_message.return_value = "display_message"
 
-        facade = StreamFacade.create_from_stream(self._stream, self._source, self._logger, self._max_workers)
+        facade = StreamFacade.create_from_stream(self._stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
 
         expected_display_message = "display_message"
         e = ExceptionWithDisplayMessage("display_message")
@@ -338,7 +362,9 @@ class StreamFacadeTest(unittest.TestCase):
 )
 def test_get_error_display_message(exception, expected_display_message):
     stream = Mock()
-    facade = StreamFacade(stream)
+    legacy_stream = Mock()
+    cursor = Mock(spec=Cursor)
+    facade = StreamFacade(stream, legacy_stream, cursor)
 
     display_message = facade.get_error_display_message(exception)
 

unit_tests/sources/streams/concurrent/test_concurrent_partition_generator.py

@@ -21,12 +21,13 @@ def test_partition_generator(slices):
 
     stream = Mock()
     message_repository = Mock()
-    partitions = [StreamPartition(stream, s, message_repository) for s in slices]
-    stream.generate.return_value = iter(partitions)
-
     sync_mode = SyncMode.full_refresh
+    cursor_field = None
+    state = None
+    partitions = [StreamPartition(stream, s, message_repository, sync_mode, cursor_field, state) for s in slices]
+    stream.generate.return_value = iter(partitions)
 
-    partition_generator.generate_partitions(stream, sync_mode)
+    partition_generator.generate_partitions(stream)
 
     actual_partitions = []
     while partition := queue.get(False):

unit_tests/sources/streams/concurrent/test_cursor.py

@@ -0,0 +1,130 @@
+from typing import Any, Mapping, Optional
+from unittest import TestCase
+from unittest.mock import Mock
+
+import pytest
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
+from airbyte_cdk.sources.message import MessageRepository
+from airbyte_cdk.sources.streams.concurrent.cursor import Comparable, ConcurrentCursor, CursorField
+from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
+from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
+
+_A_STREAM_NAME = "a stream name"
+_A_STREAM_NAMESPACE = "a stream namespace"
+_ANY_STATE = None
+_A_CURSOR_FIELD_KEY = "a_cursor_field_key"
+_NO_PARTITION_IDENTIFIER = None
+_NO_SLICE = None
+_NO_SLICE_BOUNDARIES = None
+_LOWER_SLICE_BOUNDARY_FIELD = "lower_boundary"
+_UPPER_SLICE_BOUNDARY_FIELD = "upper_boundary"
+_SLICE_BOUNDARY_FIELDS = (_LOWER_SLICE_BOUNDARY_FIELD, _UPPER_SLICE_BOUNDARY_FIELD)
+_A_VERY_HIGH_CURSOR_VALUE = 1000000000
+
+
+def _partition(_slice: Optional[Mapping[str, Any]]) -> Partition:
+    partition = Mock(spec=Partition)
+    partition.to_slice.return_value = _slice
+    return partition
+
+
+def _record(cursor_value: Comparable) -> Record:
+    return Record(data={_A_CURSOR_FIELD_KEY: cursor_value})
+
+
+class ConcurrentCursorTest(TestCase):
+    def setUp(self) -> None:
+        self._message_repository = Mock(spec=MessageRepository)
+        self._state_manager = Mock(spec=ConnectorStateManager)
+
+    def _cursor_with_slice_boundary_fields(self) -> ConcurrentCursor:
+        return ConcurrentCursor(
+            _A_STREAM_NAME,
+            _A_STREAM_NAMESPACE,
+            _ANY_STATE,
+            self._message_repository,
+            self._state_manager,
+            CursorField(_A_CURSOR_FIELD_KEY),
+            _SLICE_BOUNDARY_FIELDS,
+        )
+
+    def _cursor_without_slice_boundary_fields(self) -> ConcurrentCursor:
+        return ConcurrentCursor(
+            _A_STREAM_NAME,
+            _A_STREAM_NAMESPACE,
+            _ANY_STATE,
+            self._message_repository,
+            self._state_manager,
+            CursorField(_A_CURSOR_FIELD_KEY),
+            None,
+        )
+
+    def test_given_boundary_fields_when_close_partition_then_emit_state(self) -> None:
+        self._cursor_with_slice_boundary_fields().close_partition(
+            _partition(
+                {_LOWER_SLICE_BOUNDARY_FIELD: 12, _UPPER_SLICE_BOUNDARY_FIELD: 30},
+            )
+        )
+
+        self._message_repository.emit_message.assert_called_once_with(self._state_manager.create_state_message.return_value)
+        self._state_manager.update_state_for_stream.assert_called_once_with(
+            _A_STREAM_NAME,
+            _A_STREAM_NAMESPACE,
+            {
+                "slices": [
+                    {
+                        "start": 12,
+                        "end": 30,
+                    },
+                ]
+            },
+        )
+
+    def test_given_boundary_fields_and_record_observed_when_close_partition_then_ignore_records(self) -> None:
+        cursor = self._cursor_with_slice_boundary_fields()
+        cursor.observe(_record(_A_VERY_HIGH_CURSOR_VALUE))
+
+        cursor.close_partition(_partition({_LOWER_SLICE_BOUNDARY_FIELD: 12, _UPPER_SLICE_BOUNDARY_FIELD: 30}))
+
+        assert self._state_manager.update_state_for_stream.call_args_list[0].args[2]["slices"][0]["end"] != _A_VERY_HIGH_CURSOR_VALUE
+
+    def test_given_no_boundary_fields_when_close_partition_then_emit_state(self) -> None:
+        cursor = self._cursor_without_slice_boundary_fields()
+        cursor.observe(_record(10))
+        cursor.close_partition(_partition(_NO_SLICE))
+
+        self._state_manager.update_state_for_stream.assert_called_once_with(
+            _A_STREAM_NAME,
+            _A_STREAM_NAMESPACE,
+            {
+                "slices": [
+                    {
+                        "start": 0,
+                        "end": 10,
+                    },
+                ]
+            },
+        )
+
+    def test_given_no_boundary_fields_when_close_multiple_partitions_then_raise_exception(self) -> None:
+        cursor = self._cursor_without_slice_boundary_fields()
+        cursor.observe(_record(10))
+        cursor.close_partition(_partition(_NO_SLICE))
+
+        with pytest.raises(ValueError):
+            cursor.close_partition(_partition(_NO_SLICE))
+
+    def test_given_no_records_observed_when_close_partition_then_do_not_emit_state(self) -> None:
+        cursor = self._cursor_without_slice_boundary_fields()
+        cursor.close_partition(_partition(_NO_SLICE))
+        assert self._message_repository.emit_message.call_count == 0
+
+    def test_given_slice_boundaries_and_no_slice_when_close_partition_then_raise_error(self) -> None:
+        cursor = self._cursor_with_slice_boundary_fields()
+        with pytest.raises(KeyError):
+            cursor.close_partition(_partition(_NO_SLICE))
+
+    def test_given_slice_boundaries_not_matching_slice_when_close_partition_then_raise_error(self) -> None:
+        cursor = self._cursor_with_slice_boundary_fields()
+        with pytest.raises(KeyError):
+            cursor.close_partition(_partition({"not_matching_key": "value"}))

unit_tests/sources/streams/concurrent/test_thread_based_concurrent_stream.py

@@ -5,9 +5,9 @@
 import unittest
 from unittest.mock import Mock, call
 
-import pytest
 from airbyte_cdk.models import AirbyteStream, SyncMode
 from airbyte_cdk.sources.streams.concurrent.availability_strategy import STREAM_AVAILABLE
+from airbyte_cdk.sources.streams.concurrent.cursor import Cursor
 from airbyte_cdk.sources.streams.concurrent.partitions.partition import Partition
 from airbyte_cdk.sources.streams.concurrent.partitions.record import Record
 from airbyte_cdk.sources.streams.concurrent.thread_based_concurrent_stream import ThreadBasedConcurrentStream
@@ -25,6 +25,7 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         self._slice_logger = Mock()
         self._logger = Mock()
         self._message_repository = Mock()
+        self._cursor = Mock(spec=Cursor)
         self._stream = ThreadBasedConcurrentStream(
             self._partition_generator,
             self._max_workers,
@@ -39,6 +40,7 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
             1,
             2,
             0,
+            cursor=self._cursor,
         )
 
     def test_get_json_schema(self):
@@ -76,17 +78,20 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         with self.assertRaises(Exception):
             self._stream._check_for_errors(futures)
 
-    def test_read_raises_an_exception_if_a_partition_raises_an_exception(self):
+    def test_read_observe_records_and_close_partition(self):
         partition = Mock(spec=Partition)
-        partition.read.side_effect = RuntimeError("error")
+        expected_records = [Record({"id": 1}), Record({"id": "2"})]
+        partition.read.return_value = expected_records
+        partition.to_slice.return_value = {"slice": "slice"}
+        self._slice_logger.should_log_slice_message.return_value = False
+
         self._partition_generator.generate.return_value = [partition]
-        with pytest.raises(RuntimeError):
-            list(self._stream.read())
+        actual_records = list(self._stream.read())
+
+        assert expected_records == actual_records
 
-    def test_read_raises_an_exception_if_partition_generator_raises_an_exception(self):
-        self._partition_generator.generate.side_effect = RuntimeError("error")
-        with pytest.raises(RuntimeError):
-            list(self._stream.read())
+        self._cursor.observe.has_calls([call(record) for record in expected_records])
+        self._cursor.close_partition.assert_called_once_with(partition)
 
     def test_read_no_slice_message(self):
         partition = Mock(spec=Partition)
@@ -218,7 +223,6 @@ class ThreadBasedConcurrentStreamTest(unittest.TestCase):
         assert expected_airbyte_stream == airbyte_stream
 
     def test_as_airbyte_stream_with_a_cursor(self):
-
         json_schema = {
             "type": "object",
             "properties": {

unit_tests/sources/streams/test_stream_read.py

@@ -12,6 +12,7 @@ from airbyte_cdk.models import Type as MessageType
 from airbyte_cdk.sources.message import InMemoryMessageRepository
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
+from airbyte_cdk.sources.streams.concurrent.cursor import NoopCursor
 from airbyte_cdk.sources.streams.core import StreamData
 from airbyte_cdk.sources.utils.schema_helpers import InternalConfig
 from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger
@@ -19,6 +20,7 @@ from airbyte_cdk.sources.utils.slice_logger import DebugSliceLogger
 _A_CURSOR_FIELD = ["NESTED", "CURSOR"]
 _DEFAULT_INTERNAL_CONFIG = InternalConfig()
 _STREAM_NAME = "STREAM"
+_NO_STATE = None
 
 
 class _MockStream(Stream):
@@ -57,7 +59,7 @@ def _concurrent_stream(slice_to_partition_mapping, slice_logger, logger, message
     source = Mock()
     source._slice_logger = slice_logger
     source.message_repository = message_repository
-    stream = StreamFacade.create_from_stream(stream, source, logger, 1)
+    stream = StreamFacade.create_from_stream(stream, source, logger, 1, _NO_STATE, NoopCursor())
     stream.logger.setLevel(logger.level)
     return stream
 

unit_tests/sources/test_abstract_source.py

@@ -996,10 +996,11 @@ class TestIncrementalRead:
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 # stream 1 slice 2
                 _as_record("s1", stream_output[0]),
-                _as_record("s1", stream_output[1]),
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
+                _as_record("s1", stream_output[1]),
                 _as_record("s1", stream_output[2]),
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
+                _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 _as_stream_status("s1", AirbyteStreamStatus.COMPLETE),
                 # stream 2 slice 1
                 _as_stream_status("s2", AirbyteStreamStatus.STARTED),
@@ -1011,17 +1012,18 @@ class TestIncrementalRead:
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 # stream 2 slice 2
                 _as_record("s2", stream_output[0]),
-                _as_record("s2", stream_output[1]),
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
+                _as_record("s2", stream_output[1]),
                 _as_record("s2", stream_output[2]),
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
+                _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 _as_stream_status("s2", AirbyteStreamStatus.COMPLETE),
             ]
         )
 
         messages = _fix_emitted_at(list(src.read(logger, {}, catalog, state=input_state)))
 
-        assert expected == messages
+        assert messages == expected
 
     @pytest.mark.parametrize(
         "per_stream_enabled",
@@ -1108,11 +1110,12 @@ class TestIncrementalRead:
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 # stream 1 slice 2
                 stream_data_to_airbyte_message("s1", stream_output[0]),
+                _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 stream_data_to_airbyte_message("s1", stream_output[1]),
                 stream_data_to_airbyte_message("s1", stream_output[2]),
-                _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 stream_data_to_airbyte_message("s1", stream_output[3]),
                 _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
+                _as_state({"s1": state}, "s1", state) if per_stream_enabled else _as_state({"s1": state}),
                 _as_stream_status("s1", AirbyteStreamStatus.COMPLETE),
                 # stream 2 slice 1
                 _as_stream_status("s2", AirbyteStreamStatus.STARTED),
@@ -1125,33 +1128,33 @@ class TestIncrementalRead:
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 # stream 2 slice 2
                 stream_data_to_airbyte_message("s2", stream_output[0]),
+                _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 stream_data_to_airbyte_message("s2", stream_output[1]),
                 stream_data_to_airbyte_message("s2", stream_output[2]),
-                _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 stream_data_to_airbyte_message("s2", stream_output[3]),
                 _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
+                _as_state({"s1": state, "s2": state}, "s2", state) if per_stream_enabled else _as_state({"s1": state, "s2": state}),
                 _as_stream_status("s2", AirbyteStreamStatus.COMPLETE),
             ]
         )
 
         messages = _fix_emitted_at(list(src.read(logger, {}, catalog, state=input_state)))
 
-        assert expected == messages
+        assert messages == expected
 
 
 def test_checkpoint_state_from_stream_instance():
     teams_stream = MockStreamOverridesStateMethod()
     managers_stream = StreamNoStateMethod()
-    src = MockSource(streams=[teams_stream, managers_stream])
     state_manager = ConnectorStateManager({"teams": teams_stream, "managers": managers_stream}, [])
 
     # The stream_state passed to checkpoint_state() should be ignored since stream implements state function
     teams_stream.state = {"updated_at": "2022-09-11"}
-    actual_message = src._checkpoint_state(teams_stream, {"ignored": "state"}, state_manager)
+    actual_message = teams_stream._checkpoint_state({"ignored": "state"}, state_manager, True)
     assert actual_message == _as_state({"teams": {"updated_at": "2022-09-11"}}, "teams", {"updated_at": "2022-09-11"})
 
     # The stream_state passed to checkpoint_state() should be used since the stream does not implement state function
-    actual_message = src._checkpoint_state(managers_stream, {"updated": "expected_here"}, state_manager)
+    actual_message = managers_stream._checkpoint_state({"updated": "expected_here"}, state_manager, True)
     assert actual_message == _as_state(
         {"teams": {"updated_at": "2022-09-11"}, "managers": {"updated": "expected_here"}}, "managers", {"updated": "expected_here"}
     )