airbyte-cdk 0.53.9__py3-none-any.whl → 0.55.0__py3-none-any.whl

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

@@ -1,25 +1,23 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
 import logging
 
 from airbyte_cdk.sources.message import InMemoryMessageRepository
+from airbyte_cdk.sources.streams.concurrent.default_stream import DefaultStream
 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.utils.slice_logger import AlwaysLogSliceLogger
 from unit_tests.sources.file_based.scenarios.scenario_builder import TestScenarioBuilder
 from unit_tests.sources.streams.concurrent.scenarios.thread_based_concurrent_stream_source_builder import (
     AlwaysAvailableAvailabilityStrategy,
     ConcurrentSourceBuilder,
     InMemoryPartition,
     InMemoryPartitionGenerator,
-    NeverLogSliceLogger,
 )
 
-_id_only_stream = ThreadBasedConcurrentStream(
-    partition_generator=InMemoryPartitionGenerator([InMemoryPartition("partition1", None, [Record({"id": "1"}), Record({"id": "2"})])]),
-    max_workers=1,
+_id_only_stream = DefaultStream(
+    partition_generator=InMemoryPartitionGenerator(
+        [InMemoryPartition("partition1", "stream1", None, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")])]
+    ),
     name="stream1",
     json_schema={
         "type": "object",
@@ -30,15 +28,13 @@ _id_only_stream = ThreadBasedConcurrentStream(
     availability_strategy=AlwaysAvailableAvailabilityStrategy(),
     primary_key=[],
     cursor_field=None,
-    slice_logger=NeverLogSliceLogger(),
     logger=logging.getLogger("test_logger"),
-    message_repository=None,
-    timeout_seconds=300,
 )
 
-_id_only_stream_with_slice_logger = ThreadBasedConcurrentStream(
-    partition_generator=InMemoryPartitionGenerator([InMemoryPartition("partition1", None, [Record({"id": "1"}), Record({"id": "2"})])]),
-    max_workers=1,
+_id_only_stream_with_slice_logger = DefaultStream(
+    partition_generator=InMemoryPartitionGenerator(
+        [InMemoryPartition("partition1", "stream1", None, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")])]
+    ),
     name="stream1",
     json_schema={
         "type": "object",
@@ -49,15 +45,13 @@ _id_only_stream_with_slice_logger = ThreadBasedConcurrentStream(
     availability_strategy=AlwaysAvailableAvailabilityStrategy(),
     primary_key=[],
     cursor_field=None,
-    slice_logger=AlwaysLogSliceLogger(),
     logger=logging.getLogger("test_logger"),
-    message_repository=None,
-    timeout_seconds=300,
 )
 
-_id_only_stream_with_primary_key = ThreadBasedConcurrentStream(
-    partition_generator=InMemoryPartitionGenerator([InMemoryPartition("partition1", None, [Record({"id": "1"}), Record({"id": "2"})])]),
-    max_workers=1,
+_id_only_stream_with_primary_key = DefaultStream(
+    partition_generator=InMemoryPartitionGenerator(
+        [InMemoryPartition("partition1", "stream1", None, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")])]
+    ),
     name="stream1",
     json_schema={
         "type": "object",
@@ -68,20 +62,16 @@ _id_only_stream_with_primary_key = ThreadBasedConcurrentStream(
     availability_strategy=AlwaysAvailableAvailabilityStrategy(),
     primary_key=["id"],
     cursor_field=None,
-    slice_logger=NeverLogSliceLogger(),
     logger=logging.getLogger("test_logger"),
-    message_repository=None,
-    timeout_seconds=300,
 )
 
-_id_only_stream_multiple_partitions = ThreadBasedConcurrentStream(
+_id_only_stream_multiple_partitions = DefaultStream(
     partition_generator=InMemoryPartitionGenerator(
         [
-            InMemoryPartition("partition1", {"p": "1"}, [Record({"id": "1"}), Record({"id": "2"})]),
-            InMemoryPartition("partition2", {"p": "2"}, [Record({"id": "3"}), Record({"id": "4"})]),
+            InMemoryPartition("partition1", "stream1", {"p": "1"}, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")]),
+            InMemoryPartition("partition2", "stream1", {"p": "2"}, [Record({"id": "3"}, "stream1"), Record({"id": "4"}, "stream1")]),
         ]
     ),
-    max_workers=1,
     name="stream1",
     json_schema={
         "type": "object",
@@ -92,20 +82,16 @@ _id_only_stream_multiple_partitions = ThreadBasedConcurrentStream(
     availability_strategy=AlwaysAvailableAvailabilityStrategy(),
     primary_key=[],
     cursor_field=None,
-    slice_logger=NeverLogSliceLogger(),
     logger=logging.getLogger("test_logger"),
-    message_repository=None,
-    timeout_seconds=300,
 )
 
-_id_only_stream_multiple_partitions_concurrency_level_two = ThreadBasedConcurrentStream(
+_id_only_stream_multiple_partitions_concurrency_level_two = DefaultStream(
     partition_generator=InMemoryPartitionGenerator(
         [
-            InMemoryPartition("partition1", {"p": "1"}, [Record({"id": "1"}), Record({"id": "2"})]),
-            InMemoryPartition("partition2", {"p": "2"}, [Record({"id": "3"}), Record({"id": "4"})]),
+            InMemoryPartition("partition1", "stream1", {"p": "1"}, [Record({"id": "1"}, "stream1"), Record({"id": "2"}, "stream1")]),
+            InMemoryPartition("partition2", "stream1", {"p": "2"}, [Record({"id": "3"}, "stream1"), Record({"id": "4"}, "stream1")]),
         ]
     ),
-    max_workers=2,
     name="stream1",
     json_schema={
         "type": "object",
@@ -116,17 +102,13 @@ _id_only_stream_multiple_partitions_concurrency_level_two = ThreadBasedConcurren
     availability_strategy=AlwaysAvailableAvailabilityStrategy(),
     primary_key=[],
     cursor_field=None,
-    slice_logger=NeverLogSliceLogger(),
     logger=logging.getLogger("test_logger"),
-    message_repository=None,
-    timeout_seconds=300,
 )
 
-_stream_raising_exception = ThreadBasedConcurrentStream(
+_stream_raising_exception = DefaultStream(
     partition_generator=InMemoryPartitionGenerator(
-        [InMemoryPartition("partition1", None, [Record({"id": "1"}), ValueError("test exception")])]
+        [InMemoryPartition("partition1", "stream1", None, [Record({"id": "1"}, "stream1"), ValueError("test exception")])]
     ),
-    max_workers=1,
     name="stream1",
     json_schema={
         "type": "object",
@@ -137,10 +119,7 @@ _stream_raising_exception = ThreadBasedConcurrentStream(
     availability_strategy=AlwaysAvailableAvailabilityStrategy(),
     primary_key=[],
     cursor_field=None,
-    slice_logger=NeverLogSliceLogger(),
     logger=logging.getLogger("test_logger"),
-    message_repository=None,
-    timeout_seconds=300,
 )
 
 test_concurrent_cdk_single_stream = (
@@ -165,15 +144,14 @@ test_concurrent_cdk_single_stream = (
     .set_expected_logs(
         {
             "read": [
-                {"level": "INFO", "message": "Starting syncing ConcurrentCdkSource"},
+                {"level": "INFO", "message": "Starting syncing"},
                 {"level": "INFO", "message": "Marking stream stream1 as STARTED"},
                 {"level": "INFO", "message": "Syncing stream: stream1"},
                 {"level": "INFO", "message": "Marking stream stream1 as RUNNING"},
                 {"level": "INFO", "message": "Read 2 records from stream1 stream"},
                 {"level": "INFO", "message": "Marking stream stream1 as STOPPED"},
                 {"level": "INFO", "message": "Finished syncing stream1"},
-                {"level": "INFO", "message": "ConcurrentCdkSource runtimes"},
-                {"level": "INFO", "message": "Finished syncing ConcurrentCdkSource"},
+                {"level": "INFO", "message": "Finished syncing"},
             ]
         }
     )
@@ -202,11 +180,13 @@ test_concurrent_cdk_single_stream_with_primary_key = (
     .set_name("test_concurrent_cdk_single_stream_with_primary_key")
     .set_config({})
     .set_source_builder(
-        ConcurrentSourceBuilder().set_streams(
+        ConcurrentSourceBuilder()
+        .set_streams(
             [
                 _id_only_stream_with_primary_key,
             ]
         )
+        .set_message_repository(InMemoryMessageRepository())
     )
     .set_expected_records(
         [
@@ -239,14 +219,21 @@ test_concurrent_cdk_multiple_streams = (
     .set_name("test_concurrent_cdk_multiple_streams")
     .set_config({})
     .set_source_builder(
-        ConcurrentSourceBuilder().set_streams(
+        ConcurrentSourceBuilder()
+        .set_streams(
             [
                 _id_only_stream,
-                ThreadBasedConcurrentStream(
+                DefaultStream(
                     partition_generator=InMemoryPartitionGenerator(
-                        [InMemoryPartition("partition1", None, [Record({"id": "10", "key": "v1"}), Record({"id": "20", "key": "v2"})])]
+                        [
+                            InMemoryPartition(
+                                "partition1",
+                                "stream2",
+                                None,
+                                [Record({"id": "10", "key": "v1"}, "stream2"), Record({"id": "20", "key": "v2"}, "stream2")],
+                            )
+                        ]
                     ),
-                    max_workers=1,
                     name="stream2",
                     json_schema={
                         "type": "object",
@@ -258,13 +245,11 @@ test_concurrent_cdk_multiple_streams = (
                     availability_strategy=AlwaysAvailableAvailabilityStrategy(),
                     primary_key=[],
                     cursor_field=None,
-                    slice_logger=NeverLogSliceLogger(),
                     logger=logging.getLogger("test_logger"),
-                    message_repository=None,
-                    timeout_seconds=300,
                 ),
             ]
         )
+        .set_message_repository(InMemoryMessageRepository())
     )
     .set_expected_records(
         [
@@ -347,11 +332,13 @@ test_concurrent_cdk_single_stream_multiple_partitions = (
     .set_name("test_concurrent_cdk_single_stream_multiple_partitions")
     .set_config({})
     .set_source_builder(
-        ConcurrentSourceBuilder().set_streams(
+        ConcurrentSourceBuilder()
+        .set_streams(
             [
                 _id_only_stream_multiple_partitions,
             ]
         )
+        .set_message_repository(InMemoryMessageRepository())
     )
     .set_expected_records(
         [
@@ -385,11 +372,13 @@ test_concurrent_cdk_single_stream_multiple_partitions_concurrency_level_two = (
     .set_name("test_concurrent_cdk_single_stream_multiple_partitions_concurrency_level_2")
     .set_config({})
     .set_source_builder(
-        ConcurrentSourceBuilder().set_streams(
+        ConcurrentSourceBuilder()
+        .set_streams(
             [
                 _id_only_stream_multiple_partitions_concurrency_level_two,
             ]
         )
+        .set_message_repository(InMemoryMessageRepository())
     )
     .set_expected_records(
         [

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

@@ -1,22 +1,22 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
 import json
 import logging
 from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union
 
 from airbyte_cdk.models import ConfiguredAirbyteCatalog, ConnectorSpecification, DestinationSyncMode, SyncMode
-from airbyte_cdk.sources import AbstractSource
+from airbyte_cdk.sources.concurrent_source.concurrent_source import ConcurrentSource
+from airbyte_cdk.sources.concurrent_source.concurrent_source_adapter import ConcurrentSourceAdapter
 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.default_stream import DefaultStream
 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
@@ -37,17 +37,18 @@ class LegacyStream(Stream):
         yield from []
 
 
-class ConcurrentCdkSource(AbstractSource):
-    def __init__(self, streams: List[ThreadBasedConcurrentStream], message_repository: Optional[MessageRepository]):
+class ConcurrentCdkSource(ConcurrentSourceAdapter):
+    def __init__(self, streams: List[DefaultStream], message_repository: Optional[MessageRepository], max_workers, timeout_in_seconds):
+        concurrent_source = ConcurrentSource.create(1, 1, streams[0]._logger, NeverLogSliceLogger(), message_repository)
+        super().__init__(concurrent_source)
         self._streams = streams
-        self._message_repository = message_repository
 
     def check_connection(self, logger: logging.Logger, config: Mapping[str, Any]) -> Tuple[bool, Optional[Any]]:
         # Check is not verified because it is up to the source to implement this method
         return True, None
 
     def streams(self, config: Mapping[str, Any]) -> List[Stream]:
-        return [StreamFacade(s, LegacyStream(), NoopCursor()) for s in self._streams]
+        return [StreamFacade(s, LegacyStream(), NoopCursor(), NeverLogSliceLogger(), s._logger) for s in self._streams]
 
     def spec(self, *args: Any, **kwargs: Any) -> ConnectorSpecification:
         return ConnectorSpecification(connectionSpecification={})
@@ -56,7 +57,7 @@ class ConcurrentCdkSource(AbstractSource):
         return ConfiguredAirbyteCatalog(
             streams=[
                 ConfiguredAirbyteStream(
-                    stream=StreamFacade(s, LegacyStream(), NoopCursor()).as_airbyte_stream(),
+                    stream=StreamFacade(s, LegacyStream(), NoopCursor(), NeverLogSliceLogger(), s._logger).as_airbyte_stream(),
                     sync_mode=SyncMode.full_refresh,
                     destination_sync_mode=DestinationSyncMode.overwrite,
                 )
@@ -78,10 +79,15 @@ class InMemoryPartitionGenerator(PartitionGenerator):
 
 
 class InMemoryPartition(Partition):
-    def __init__(self, name, _slice, records):
+    def stream_name(self) -> str:
+        return self._stream_name
+
+    def __init__(self, name, stream_name, _slice, records):
         self._name = name
+        self._stream_name = stream_name
         self._slice = _slice
         self._records = records
+        self._is_closed = False
 
     def read(self) -> Iterable[Record]:
         for record_or_exception in self._records:
@@ -101,19 +107,22 @@ class InMemoryPartition(Partition):
         else:
             return hash(self._name)
 
+    def close(self) -> None:
+        self._is_closed = True
+
+    def is_closed(self) -> bool:
+        return self._is_closed
+
 
 class ConcurrentSourceBuilder(SourceBuilder[ConcurrentCdkSource]):
     def __init__(self):
-        self._streams: List[ThreadBasedConcurrentStream] = []
+        self._streams: List[DefaultStream] = []
         self._message_repository = None
 
     def build(self, configured_catalog: Optional[Mapping[str, Any]]) -> ConcurrentCdkSource:
-        for stream in self._streams:
-            if not stream._message_repository:
-                stream._message_repository = self._message_repository
-        return ConcurrentCdkSource(self._streams, self._message_repository)
+        return ConcurrentCdkSource(self._streams, self._message_repository, 1, 1)
 
-    def set_streams(self, streams: List[ThreadBasedConcurrentStream]) -> "ConcurrentSourceBuilder":
+    def set_streams(self, streams: List[DefaultStream]) -> "ConcurrentSourceBuilder":
         self._streams = streams
         return self
 

unit_tests/sources/streams/concurrent/test_adapters.py

@@ -1,7 +1,7 @@
 #
 # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
 #
-
+import logging
 import unittest
 from unittest.mock import Mock
 
@@ -17,15 +17,18 @@ 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.cursor import Cursor
 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.slice_logger import SliceLogger
 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"]
+_STREAM_NAME = "stream"
+_ANY_CURSOR = Mock(spec=Cursor)
 
 
 @pytest.mark.parametrize(
@@ -77,7 +80,7 @@ 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, _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE)
+    partition_generator = StreamPartitionGenerator(stream, message_repository, _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE, _ANY_CURSOR)
 
     partitions = list(partition_generator.generate())
     slices = [partition.to_slice() for partition in partitions]
@@ -88,16 +91,21 @@ def test_stream_partition_generator(sync_mode):
 @pytest.mark.parametrize(
     "transformer, expected_records",
     [
-        pytest.param(TypeTransformer(TransformConfig.NoTransform), [Record({"data": "1"}), Record({"data": "2"})], id="test_no_transform"),
+        pytest.param(
+            TypeTransformer(TransformConfig.NoTransform),
+            [Record({"data": "1"}, _STREAM_NAME), Record({"data": "2"}, _STREAM_NAME)],
+            id="test_no_transform",
+        ),
         pytest.param(
             TypeTransformer(TransformConfig.DefaultSchemaNormalization),
-            [Record({"data": 1}), Record({"data": 2})],
+            [Record({"data": 1}, _STREAM_NAME), Record({"data": 2}, _STREAM_NAME)],
             id="test_default_transform",
         ),
     ],
 )
 def test_stream_partition(transformer, expected_records):
     stream = Mock()
+    stream.name = _STREAM_NAME
     stream.get_json_schema.return_value = {"type": "object", "properties": {"data": {"type": ["integer"]}}}
     stream.transformer = transformer
     message_repository = InMemoryMessageRepository()
@@ -105,7 +113,7 @@ def test_stream_partition(transformer, expected_records):
     sync_mode = SyncMode.full_refresh
     cursor_field = None
     state = None
-    partition = StreamPartition(stream, _slice, message_repository, sync_mode, cursor_field, state)
+    partition = StreamPartition(stream, _slice, message_repository, sync_mode, cursor_field, state, _ANY_CURSOR)
 
     a_log_message = AirbyteMessage(
         type=MessageType.LOG,
@@ -139,7 +147,7 @@ def test_stream_partition_raising_exception(exception_type, expected_display_mes
     message_repository = InMemoryMessageRepository()
     _slice = None
 
-    partition = StreamPartition(stream, _slice, message_repository, _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE)
+    partition = StreamPartition(stream, _slice, message_repository, _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE, _ANY_CURSOR)
 
     stream.read_records.side_effect = Exception()
 
@@ -159,7 +167,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(), _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE)
+    partition = StreamPartition(stream, _slice, Mock(), _ANY_SYNC_MODE, _ANY_CURSOR_FIELD, _ANY_STATE, _ANY_CURSOR)
 
     _hash = partition.__hash__()
     assert _hash == expected_hash
@@ -176,10 +184,11 @@ class StreamFacadeTest(unittest.TestCase):
         )
         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._slice_logger = Mock()
+        self._slice_logger.should_log_slice_message.return_value = False
+        self._facade = StreamFacade(self._abstract_stream, self._legacy_stream, self._cursor, self._slice_logger, self._logger)
         self._source = Mock()
-        self._max_workers = 10
 
         self._stream = Mock()
         self._stream.primary_key = "id"
@@ -206,12 +215,16 @@ class StreamFacadeTest(unittest.TestCase):
 
     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
+            StreamFacade(
+                self._abstract_stream, self._legacy_stream, _ANY_CURSOR, Mock(spec=SliceLogger), Mock(spec=logging.Logger)
+            ).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
+        assert StreamFacade(
+            self._abstract_stream, self._legacy_stream, Mock(spec=Cursor), Mock(spec=SliceLogger), Mock(spec=logging.Logger)
+        ).supports_incremental
 
     def test_check_availability_is_delegated_to_wrapped_stream(self):
         availability = StreamAvailable()
@@ -221,8 +234,11 @@ class StreamFacadeTest(unittest.TestCase):
 
     def test_full_refresh(self):
         expected_stream_data = [{"data": 1}, {"data": 2}]
-        records = [Record(data) for data in expected_stream_data]
-        self._abstract_stream.read.return_value = records
+        records = [Record(data, "stream") for data in expected_stream_data]
+
+        partition = Mock()
+        partition.read.return_value = records
+        self._abstract_stream.generate_partitions.return_value = [partition]
 
         actual_stream_data = list(self._facade.read_records(SyncMode.full_refresh, None, None, None))
 
@@ -230,8 +246,10 @@ class StreamFacadeTest(unittest.TestCase):
 
     def test_read_records_full_refresh(self):
         expected_stream_data = [{"data": 1}, {"data": 2}]
-        records = [Record(data) for data in expected_stream_data]
-        self._abstract_stream.read.return_value = records
+        records = [Record(data, "stream") for data in expected_stream_data]
+        partition = Mock()
+        partition.read.return_value = records
+        self._abstract_stream.generate_partitions.return_value = [partition]
 
         actual_stream_data = list(self._facade.read_full_refresh(None, None, None))
 
@@ -239,8 +257,10 @@ class StreamFacadeTest(unittest.TestCase):
 
     def test_read_records_incremental(self):
         expected_stream_data = [{"data": 1}, {"data": 2}]
-        records = [Record(data) for data in expected_stream_data]
-        self._abstract_stream.read.return_value = records
+        records = [Record(data, "stream") for data in expected_stream_data]
+        partition = Mock()
+        partition.read.return_value = records
+        self._abstract_stream.generate_partitions.return_value = [partition]
 
         actual_stream_data = list(self._facade.read_incremental(None, None, None, None, None, None, None))
 
@@ -252,7 +272,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, _ANY_STATE, self._cursor)
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
 
         assert facade.name == "stream"
         assert facade.cursor_field == "cursor"
@@ -264,8 +284,8 @@ class StreamFacadeTest(unittest.TestCase):
         stream.primary_key = None
         stream.cursor_field = []
 
-        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
-        facade._abstract_stream._primary_key is None
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+        assert facade._abstract_stream._primary_key == []
 
     def test_create_from_stream_with_composite_primary_key(self):
         stream = Mock()
@@ -273,15 +293,15 @@ 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, _ANY_STATE, self._cursor)
-        facade._abstract_stream._primary_key == ["id", "name"]
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
+        assert facade._abstract_stream._primary_key == ["id", "name"]
 
     def test_create_from_stream_with_empty_list_cursor(self):
         stream = Mock()
         stream.primary_key = "id"
         stream.cursor_field = []
 
-        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, self._max_workers, _ANY_STATE, self._cursor)
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
 
         assert facade.cursor_field == []
 
@@ -291,7 +311,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, _ANY_STATE, self._cursor)
+            StreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
 
     def test_create_from_stream_raises_exception_if_primary_key_has_invalid_type(self):
         stream = Mock()
@@ -299,7 +319,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, _ANY_STATE, self._cursor)
+            StreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
 
     def test_create_from_stream_raises_exception_if_cursor_field_is_nested(self):
         stream = Mock()
@@ -308,7 +328,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, _ANY_STATE, self._cursor)
+            StreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
 
     def test_create_from_stream_with_cursor_field_as_list(self):
         stream = Mock()
@@ -316,7 +336,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, _ANY_STATE, self._cursor)
+        facade = StreamFacade.create_from_stream(stream, self._source, self._logger, _ANY_STATE, self._cursor)
         assert facade.cursor_field == "cursor"
 
     def test_create_from_stream_none_message_repository(self):
@@ -326,12 +346,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, {}, self._cursor)
+            StreamFacade.create_from_stream(self._stream, self._source, self._logger, {}, 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, _ANY_STATE, self._cursor)
+        facade = StreamFacade.create_from_stream(self._stream, self._source, self._logger, _ANY_STATE, self._cursor)
 
         expected_display_message = None
         e = Exception()
@@ -343,7 +363,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, _ANY_STATE, self._cursor)
+        facade = StreamFacade.create_from_stream(self._stream, self._source, self._logger, _ANY_STATE, self._cursor)
 
         expected_display_message = "display_message"
         e = ExceptionWithDisplayMessage("display_message")
@@ -364,7 +384,7 @@ def test_get_error_display_message(exception, expected_display_message):
     stream = Mock()
     legacy_stream = Mock()
     cursor = Mock(spec=Cursor)
-    facade = StreamFacade(stream, legacy_stream, cursor)
+    facade = StreamFacade(stream, legacy_stream, cursor, Mock().Mock(), Mock())
 
     display_message = facade.get_error_display_message(exception)
 

unit_tests/sources/streams/concurrent/test_concurrent_partition_generator.py

@@ -7,9 +7,9 @@ from unittest.mock import Mock
 
 import pytest
 from airbyte_cdk.models import SyncMode
+from airbyte_cdk.sources.concurrent_source.partition_generation_completed_sentinel import PartitionGenerationCompletedSentinel
 from airbyte_cdk.sources.streams.concurrent.adapters import StreamPartition
 from airbyte_cdk.sources.streams.concurrent.partition_enqueuer import PartitionEnqueuer
-from airbyte_cdk.sources.streams.concurrent.partitions.types import PARTITIONS_GENERATED_SENTINEL
 
 
 @pytest.mark.parametrize(
@@ -17,21 +17,22 @@ from airbyte_cdk.sources.streams.concurrent.partitions.types import PARTITIONS_G
 )
 def test_partition_generator(slices):
     queue = Queue()
-    partition_generator = PartitionEnqueuer(queue, PARTITIONS_GENERATED_SENTINEL)
+    partition_generator = PartitionEnqueuer(queue)
 
     stream = Mock()
     message_repository = Mock()
     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)
+    cursor = Mock()
+    partitions = [StreamPartition(stream, s, message_repository, sync_mode, cursor_field, state, cursor) for s in slices]
+    stream.generate_partitions.return_value = iter(partitions)
 
     partition_generator.generate_partitions(stream)
 
     actual_partitions = []
     while partition := queue.get(False):
-        if partition == PARTITIONS_GENERATED_SENTINEL:
+        if isinstance(partition, PartitionGenerationCompletedSentinel):
             break
         actual_partitions.append(partition)