Skip to content

Commit 289705d

Browse files
maxi297octavia-squidington-iii
andauthored
fix: Ensure page token 0 is being passed (#482)
Co-authored-by: octavia-squidington-iii <contact@airbyte.com>
1 parent fa7f2e6 commit 289705d

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
@@ -364,7 +364,7 @@ def _read_pages(
364364
pagination_complete = False
365365
initial_token = self._paginator.get_initial_token()
366366
next_page_token: Optional[Mapping[str, Any]] = (
367-
{"next_page_token": initial_token} if initial_token else None
367+
{"next_page_token": initial_token} if initial_token is not None else None
368368
)
369369
while not pagination_complete:
370370
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
@@ -810,6 +810,42 @@ def retriever_read_pages(_, __, ___):
810810
cursor.close_slice.assert_called_once_with(stream_slice, None)
811811

812812

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

0 commit comments

Comments
 (0)