Skip to content

Commit 09e68cd

Browse files
authored
Merge branch 'main' into devin/1744663914-pyupgrade-rules
2 parents 62ed774 + 289705d commit 09e68cd

File tree

2 files changed

+37
-1
lines changed

2 files changed

+37
-1
lines changed

airbyte_cdk/sources/declarative/retrievers/simple_retriever.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -361,7 +361,7 @@ def _read_pages(
361361
pagination_complete = False
362362
initial_token = self._paginator.get_initial_token()
363363
next_page_token: Mapping[str, Any] | None = (
364-
{"next_page_token": initial_token} if initial_token else None
364+
{"next_page_token": initial_token} if initial_token is not None else None
365365
)
366366
while not pagination_complete:
367367
response = self._fetch_next_page(stream_state, stream_slice, next_page_token)

unit_tests/sources/declarative/retrievers/test_simple_retriever.py

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -811,6 +811,42 @@ def retriever_read_pages(_, __, ___):
811811
cursor.close_slice.assert_called_once_with(stream_slice, None)
812812

813813

814+
def test_given_initial_token_is_zero_when_read_records_then_pass_initial_token():
815+
record_selector = MagicMock()
816+
record_selector.select_records.return_value = []
817+
cursor = MagicMock(spec=DeclarativeCursor)
818+
paginator = MagicMock()
819+
paginator.get_initial_token.return_value = 0
820+
paginator.next_page_token.return_value = None
821+
822+
retriever = SimpleRetriever(
823+
name="stream_name",
824+
primary_key=primary_key,
825+
requester=MagicMock(),
826+
paginator=paginator,
827+
record_selector=record_selector,
828+
stream_slicer=cursor,
829+
cursor=cursor,
830+
parameters={},
831+
config={},
832+
)
833+
stream_slice = StreamSlice(cursor_slice={}, partition={})
834+
835+
response = requests.Response()
836+
response.status_code = 200
837+
response._content = "{}".encode()
838+
839+
with patch.object(
840+
SimpleRetriever,
841+
"_fetch_next_page",
842+
return_value=response,
843+
) as fetch_next_page_mock:
844+
list(retriever.read_records(stream_slice=stream_slice, records_schema={}))
845+
fetch_next_page_mock.assert_called_once_with(
846+
cursor.get_stream_state(), stream_slice, {"next_page_token": 0}
847+
)
848+
849+
814850
def _generate_slices(number_of_slices):
815851
return [{"date": f"2022-01-0{day + 1}"} for day in range(number_of_slices)]
816852

0 commit comments

Comments
 (0)