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

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

@@ -6,27 +6,61 @@ from typing import Any, List, Mapping, Optional, Tuple, Union
 
 from airbyte_cdk.models import ConfiguredAirbyteCatalog, ConnectorSpecification, DestinationSyncMode, SyncMode
 from airbyte_cdk.sources import AbstractSource
+from airbyte_cdk.sources.connector_state_manager import ConnectorStateManager
 from airbyte_cdk.sources.message import InMemoryMessageRepository, MessageRepository
 from airbyte_cdk.sources.streams import Stream
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamFacade
+from airbyte_cdk.sources.streams.concurrent.cursor import ConcurrentCursor, CursorField, NoopCursor
 from airbyte_protocol.models import ConfiguredAirbyteStream
 from unit_tests.sources.file_based.scenarios.scenario_builder import SourceBuilder
 
+_NO_STATE = None
+
 
 class StreamFacadeSource(AbstractSource):
-    def __init__(self, streams: List[Stream], max_workers: int):
+    def __init__(
+        self,
+        streams: List[Stream],
+        max_workers: int,
+        cursor_field: Optional[CursorField] = None,
+        cursor_boundaries: Optional[Tuple[str, str]] = None,
+    ):
         self._streams = streams
         self._max_workers = max_workers
+        self._message_repository = InMemoryMessageRepository()
+        self._cursor_field = cursor_field
+        self._cursor_boundaries = cursor_boundaries
 
     def check_connection(self, logger: logging.Logger, config: Mapping[str, Any]) -> Tuple[bool, Optional[Any]]:
         return True, None
 
     def streams(self, config: Mapping[str, Any]) -> List[Stream]:
-        return [StreamFacade.create_from_stream(stream, self, stream.logger, self._max_workers) for stream in self._streams]
+        state_manager = ConnectorStateManager(stream_instance_map={s.name: s for s in self._streams}, state=_NO_STATE)
+        return [
+            StreamFacade.create_from_stream(
+                stream,
+                self,
+                stream.logger,
+                self._max_workers,
+                _NO_STATE,
+                ConcurrentCursor(
+                    stream.name,
+                    stream.namespace,
+                    _NO_STATE,
+                    self.message_repository,  # type: ignore  # for this source specifically, we always return `InMemoryMessageRepository`
+                    state_manager,
+                    self._cursor_field,
+                    self._cursor_boundaries,
+                )
+                if self._cursor_field
+                else NoopCursor(),
+            )
+            for stream in self._streams
+        ]
 
     @property
     def message_repository(self) -> Union[None, MessageRepository]:
-        return InMemoryMessageRepository()
+        return self._message_repository
 
     def spec(self, logger: logging.Logger) -> ConnectorSpecification:
         return ConnectorSpecification(connectionSpecification={})
@@ -49,14 +83,21 @@ class StreamFacadeSourceBuilder(SourceBuilder[StreamFacadeSource]):
         self._source = None
         self._streams = []
         self._max_workers = 1
+        self._cursor_field = None
+        self._cursor_boundaries = None
 
     def set_streams(self, streams: List[Stream]) -> "StreamFacadeSourceBuilder":
         self._streams = streams
         return self
 
-    def set_max_workers(self, max_workers: int):
+    def set_max_workers(self, max_workers: int) -> "StreamFacadeSourceBuilder":
         self._max_workers = max_workers
         return self
 
+    def set_incremental(self, cursor_field: CursorField, cursor_boundaries: Optional[Tuple[str, str]]) -> "StreamFacadeSourceBuilder":
+        self._cursor_field = cursor_field
+        self._cursor_boundaries = cursor_boundaries
+        return self
+
     def build(self, configured_catalog: Optional[Mapping[str, Any]]) -> StreamFacadeSource:
-        return StreamFacadeSource(self._streams, self._max_workers)
+        return StreamFacadeSource(self._streams, self._max_workers, self._cursor_field, self._cursor_boundaries)

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

@@ -1,26 +1,25 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-from typing import Any, Iterable, List, Mapping, Optional, Union
+from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union
 
 from airbyte_cdk.models import SyncMode
 from airbyte_cdk.sources.streams import Stream
+from airbyte_cdk.sources.streams.concurrent.cursor import CursorField
 from airbyte_cdk.sources.streams.core import StreamData
-from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenarioBuilder
+from unit_tests.sources.file_based.scenarios.scenario_builder import IncrementalScenarioConfig, TestScenarioBuilder
 from unit_tests.sources.streams.concurrent.scenarios.stream_facade_builder import StreamFacadeSourceBuilder
 
 
 class _MockStream(Stream):
     def __init__(
         self,
-        slice_key,
-        slice_values_to_records_or_exception: Mapping[Optional[str], List[Union[Mapping[str, Any], Exception]]],
+        slices_and_records_or_exception: Iterable[Tuple[Optional[Mapping[str, Any]], Iterable[Union[Exception, Mapping[str, Any]]]]],
         name,
         json_schema,
         primary_key=None,
     ):
-        self._slice_key = slice_key
-        self._slice_values_to_records = slice_values_to_records_or_exception
+        self._slices_and_records_or_exception = slices_and_records_or_exception
         self._name = name
         self._json_schema = json_schema
         self._primary_key = primary_key
@@ -32,19 +31,12 @@ class _MockStream(Stream):
         stream_slice: Optional[Mapping[str, Any]] = None,
         stream_state: Optional[Mapping[str, Any]] = None,
     ) -> Iterable[StreamData]:
-        for record_or_exception in self._get_record_or_exception_iterable(stream_slice):
-            if isinstance(record_or_exception, Exception):
-                raise record_or_exception
-            else:
-                yield record_or_exception
-
-    def _get_record_or_exception_iterable(
-        self, stream_slice: Optional[Mapping[str, Any]] = None
-    ) -> Iterable[Union[Mapping[str, Any], Exception]]:
-        if stream_slice is None:
-            return self._slice_values_to_records[None]
-        else:
-            return self._slice_values_to_records[stream_slice[self._slice_key]]
+        for _slice, records_or_exception in self._slices_and_records_or_exception:
+            if stream_slice == _slice:
+                for item in records_or_exception:
+                    if isinstance(item, Exception):
+                        raise item
+                    yield item
 
     @property
     def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
@@ -60,16 +52,16 @@ class _MockStream(Stream):
     def stream_slices(
         self, *, sync_mode: SyncMode, cursor_field: Optional[List[str]] = None, stream_state: Optional[Mapping[str, Any]] = None
     ) -> Iterable[Optional[Mapping[str, Any]]]:
-        if self._slice_key:
-            for slice_value in self._slice_values_to_records.keys():
-                yield {self._slice_key: slice_value}
+        if self._slices_and_records_or_exception:
+            yield from [_slice for _slice, records_or_exception in self._slices_and_records_or_exception]
         else:
             yield None
 
 
 _stream1 = _MockStream(
-    None,
-    {None: [{"id": "1"}, {"id": "2"}]},
+    [
+        (None, [{"id": "1"}, {"id": "2"}]),
+    ],
     "stream1",
     json_schema={
         "type": "object",
@@ -80,8 +72,9 @@ _stream1 = _MockStream(
 )
 
 _stream_raising_exception = _MockStream(
-    None,
-    {None: [{"id": "1"}, ValueError("test exception")]},
+    [
+        (None, [{"id": "1"}, ValueError("test exception")]),
+    ],
     "stream1",
     json_schema={
         "type": "object",
@@ -92,8 +85,9 @@ _stream_raising_exception = _MockStream(
 )
 
 _stream_with_primary_key = _MockStream(
-    None,
-    {None: [{"id": "1"}, {"id": "2"}]},
+    [
+        (None, [{"id": "1"}, {"id": "2"}]),
+    ],
     "stream1",
     json_schema={
         "type": "object",
@@ -105,8 +99,9 @@ _stream_with_primary_key = _MockStream(
 )
 
 _stream2 = _MockStream(
-    None,
-    {None: [{"id": "A"}, {"id": "B"}]},
+    [
+        (None, [{"id": "A"}, {"id": "B"}]),
+    ],
     "stream2",
     json_schema={
         "type": "object",
@@ -117,8 +112,9 @@ _stream2 = _MockStream(
 )
 
 _stream_with_single_slice = _MockStream(
-    "slice_key",
-    {"s1": [{"id": "1"}, {"id": "2"}]},
+    [
+        ({"slice_key": "s1"}, [{"id": "1"}, {"id": "2"}]),
+    ],
     "stream1",
     json_schema={
         "type": "object",
@@ -129,11 +125,10 @@ _stream_with_single_slice = _MockStream(
 )
 
 _stream_with_multiple_slices = _MockStream(
-    "slice_key",
-    {
-        "s1": [{"id": "1"}, {"id": "2"}],
-        "s2": [{"id": "3"}, {"id": "4"}],
-    },
+    [
+        ({"slice_key": "s1"}, [{"id": "1"}, {"id": "2"}]),
+        ({"slice_key": "s2"}, [{"id": "3"}, {"id": "4"}]),
+    ],
     "stream1",
     json_schema={
         "type": "object",
@@ -384,3 +379,125 @@ test_stream_facade_single_stream_with_multiple_slices_with_concurrency_level_two
     )
     .build()
 )
+
+
+test_incremental_stream_with_slice_boundaries = (
+    TestScenarioBuilder()
+    .set_name("test_incremental_stream_with_slice_boundaries")
+    .set_config({})
+    .set_source_builder(
+        StreamFacadeSourceBuilder()
+        .set_streams(
+            [
+                _MockStream(
+                    [
+                        ({"from": 0, "to": 1}, [{"id": "1", "cursor_field": 0}, {"id": "2", "cursor_field": 1}]),
+                        ({"from": 1, "to": 2}, [{"id": "3", "cursor_field": 2}, {"id": "4", "cursor_field": 3}]),
+                    ],
+                    "stream1",
+                    json_schema={
+                        "type": "object",
+                        "properties": {
+                            "id": {"type": ["null", "string"]},
+                        },
+                    },
+                )
+            ]
+        )
+        .set_incremental(CursorField("cursor_field"), ("from", "to"))
+    )
+    .set_expected_records(
+        [
+            {"data": {"id": "1", "cursor_field": 0}, "stream": "stream1"},
+            {"data": {"id": "2", "cursor_field": 1}, "stream": "stream1"},
+            {"stream1": {"slices": [{"start": 0, "end": 1}]}},
+            {"data": {"id": "3", "cursor_field": 2}, "stream": "stream1"},
+            {"data": {"id": "4", "cursor_field": 3}, "stream": "stream1"},
+            {"stream1": {"slices": [{"start": 0, "end": 1}, {"start": 1, "end": 2}]}},
+        ]
+    )
+    .set_log_levels({"ERROR", "WARN", "WARNING", "INFO", "DEBUG"})
+    .set_incremental_scenario_config(
+        IncrementalScenarioConfig(
+            input_state=[],
+        )
+    )
+    .build()
+)
+
+
+_NO_SLICE_BOUNDARIES = None
+test_incremental_stream_without_slice_boundaries = (
+    TestScenarioBuilder()
+    .set_name("test_incremental_stream_without_slice_boundaries")
+    .set_config({})
+    .set_source_builder(
+        StreamFacadeSourceBuilder()
+        .set_streams(
+            [
+                _MockStream(
+                    [
+                        (None, [{"id": "1", "cursor_field": 0}, {"id": "2", "cursor_field": 3}]),
+                    ],
+                    "stream1",
+                    json_schema={
+                        "type": "object",
+                        "properties": {
+                            "id": {"type": ["null", "string"]},
+                        },
+                    },
+                )
+            ]
+        )
+        .set_incremental(CursorField("cursor_field"), _NO_SLICE_BOUNDARIES)
+    )
+    .set_expected_records(
+        [
+            {"data": {"id": "1", "cursor_field": 0}, "stream": "stream1"},
+            {"data": {"id": "2", "cursor_field": 3}, "stream": "stream1"},
+            {"stream1": {"slices": [{"start": 0, "end": 3}]}},
+        ]
+    )
+    .set_log_levels({"ERROR", "WARN", "WARNING", "INFO", "DEBUG"})
+    .set_incremental_scenario_config(
+        IncrementalScenarioConfig(
+            input_state=[],
+        )
+    )
+    .build()
+)
+
+test_incremental_stream_with_many_slices_but_without_slice_boundaries = (
+    TestScenarioBuilder()
+    .set_name("test_incremental_stream_with_many_slices_byt_without_slice_boundaries")
+    .set_config({})
+    .set_source_builder(
+        StreamFacadeSourceBuilder()
+        .set_streams(
+            [
+                _MockStream(
+                    [
+                        ({"parent_id": 1}, [{"id": "1", "cursor_field": 0}]),
+                        ({"parent_id": 309}, [{"id": "3", "cursor_field": 0}]),
+                    ],
+                    "stream1",
+                    json_schema={
+                        "type": "object",
+                        "properties": {
+                            "id": {"type": ["null", "string"]},
+                        },
+                    },
+                )
+            ]
+        )
+        .set_incremental(CursorField("cursor_field"), _NO_SLICE_BOUNDARIES)
+    )
+    .set_expected_read_error(ValueError, "test exception")
+    .set_log_levels({"ERROR", "WARN", "WARNING", "INFO", "DEBUG"})
+    .set_incremental_scenario_config(
+        IncrementalScenarioConfig(
+            input_state=[],
+        )
+    )
+    .build()
+)

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):