ref: specify stream_slice check

Signed-off-by: Artem Inzhyyants <artem.inzhyyants@gmail.com>
airbytehq · Jan 30, 2025 · 0466cc7 · 0466cc7
1 parent c87671e
commit 0466cc7
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 7 deletions.
diff --git a/airbyte_cdk/sources/declarative/declarative_stream.py b/airbyte_cdk/sources/declarative/declarative_stream.py
@@ -138,7 +138,9 @@ def read_records(
         """
         :param: stream_state We knowingly avoid using stream_state as we want cursors to manage their own state.
         """
-        if stream_slice is None or stream_slice == {}:
+        if stream_slice is None or (
+            not isinstance(stream_slice, StreamSlice) and stream_slice == {}
+        ):
             # As the parameter is Optional, many would just call `read_records(sync_mode)` during testing without specifying the field
             # As part of the declarative model without custom components, this should never happen as the CDK would wire up a
             # SinglePartitionRouter that would create this StreamSlice properly

diff --git a/unit_tests/sources/declarative/async_job/test_integration.py b/unit_tests/sources/declarative/async_job/test_integration.py
@@ -36,11 +36,6 @@
 _NO_LIMIT = 10000
 
 
-class MockAsyncJobPartitionRouter(AsyncJobPartitionRouter):
-    def fetch_records(self, async_jobs: Iterable[AsyncJob]) -> Iterable[Mapping[str, Any]]:
-        yield from [{"record_field": 10}]
-
-
 class MockAsyncJobRepository(AsyncJobRepository):
     def start(self, stream_slice: StreamSlice) -> AsyncJob:
         return AsyncJob("a_job_id", stream_slice)
@@ -87,7 +82,7 @@ def streams(self, config: Mapping[str, Any]) -> List[Stream]:
                     config={},
                     parameters={},
                     record_selector=noop_record_selector,
-                    stream_slicer=MockAsyncJobPartitionRouter(
+                    stream_slicer=AsyncJobPartitionRouter(
                         stream_slicer=self._stream_slicer,
                         job_orchestrator_factory=lambda stream_slices: AsyncJobOrchestrator(
                             MockAsyncJobRepository(),