From ee33548508b58fe4723f2c93eb5ac4ab301d20fe Mon Sep 17 00:00:00 2001 From: Anton Burnashev Date: Wed, 3 Apr 2024 10:39:46 +0300 Subject: [PATCH] Ensure total count is an integer in OffsetPaginator (#1172) --- dlt/sources/helpers/rest_client/paginators.py | 8 ++++++++ .../sources/helpers/rest_client/test_paginators.py | 13 +++++++++++++ 2 files changed, 21 insertions(+) diff --git a/dlt/sources/helpers/rest_client/paginators.py b/dlt/sources/helpers/rest_client/paginators.py index 48dfdf6e4f..ce414322a0 100644 --- a/dlt/sources/helpers/rest_client/paginators.py +++ b/dlt/sources/helpers/rest_client/paginators.py @@ -88,6 +88,14 @@ def update_state(self, response: Response) -> None: f"Total count not found in response for {self.__class__.__name__}" ) + try: + total = int(total) + except ValueError: + raise ValueError( + f"Total count is not an integer in response for {self.__class__.__name__}. " + f"Expected an integer, got {total}" + ) + self.offset += self.limit if self.offset >= total: diff --git a/tests/sources/helpers/rest_client/test_paginators.py b/tests/sources/helpers/rest_client/test_paginators.py index bd38a2e421..4d086f1486 100644 --- a/tests/sources/helpers/rest_client/test_paginators.py +++ b/tests/sources/helpers/rest_client/test_paginators.py @@ -186,6 +186,19 @@ def test_update_state(self): paginator.update_state(response) assert paginator.has_next_page is False + def test_update_state_with_string_total(self): + paginator = OffsetPaginator(0, 10) + response = Mock(Response, json=lambda: {"total": "20"}) + paginator.update_state(response) + assert paginator.offset == 10 + assert paginator.has_next_page is True + + def test_update_state_with_invalid_total(self): + paginator = OffsetPaginator(0, 10) + response = Mock(Response, json=lambda: {"total": "invalid"}) + with pytest.raises(ValueError): + paginator.update_state(response) + def test_update_state_without_total(self): paginator = OffsetPaginator(0, 10) response = Mock(Response, json=lambda: {})