create scroll method

flareio/api_client.py

@@ -145,3 +145,54 @@ def delete(
             json=json,
             headers=headers,
         )
+
+    def scroll(
+        self,
+        *,
+        method: t.Union[
+            t.Literal["GET"],
+            t.Literal["POST"],
+        ],
+        url: str,
+        params: t.Optional[t.Dict[str, t.Any]] = None,
+        json: t.Optional[t.Dict[str, t.Any]] = None,
+        headers: t.Optional[t.Dict[str, t.Any]] = None,
+    ) -> t.Iterator[requests.Response]:
+        if method not in {"GET", "POST"}:
+            raise Exception("Scrolling is only supported for GET or POST")
+
+        from_in_params: bool = "from" in (params or {})
+        from_in_json: bool = "from" in (json or {})
+
+        if not (from_in_params or from_in_json):
+            raise Exception("You must specify from eithier in params or in json")
+        if from_in_params and from_in_json:
+            raise Exception("You can't specify from both in params and in json")
+
+        while True:
+            resp = self._request(
+                method=method,
+                url=url,
+                params=params,
+                json=json,
+                headers=headers,
+            )
+            resp.raise_for_status()
+
+            yield resp
+
+            resp_body = resp.json()
+
+            if "next" not in resp_body:
+                raise Exception(
+                    "'next' was not found in the response body. Are you sure it supports scrolling?"
+                )
+
+            next_page: t.Optional[str] = resp_body["next"]
+            if not next_page:
+                break
+
+            if params and from_in_params:
+                params["from"] = next_page
+            if json and from_in_json:
+                json["from"] = next_page

tests/test_api_client.py

@@ -1,4 +1,5 @@
 import pytest
+import requests
 import requests_mock
 
 from datetime import datetime
@@ -163,3 +164,76 @@ def test_wrapped_methods() -> None:
         client.delete("https://api.flare.io/hello-delete")
         assert mocker.last_request.url == "https://api.flare.io/hello-delete"
         assert mocker.last_request.headers["Authorization"] == "Bearer test-token-hello"
+
+
+def test_scroll() -> None:
+    api_client = _get_test_client()
+
+    # This should make no http call.
+    with requests_mock.Mocker() as mocker:
+        response_iterator: t.Iterator[requests.Response] = api_client.scroll(
+            method="GET",
+            url="https://api.flare.io/leaksdb/v2/sources",
+            params={
+                "from": None,
+            },
+        )
+        assert len(mocker.request_history) == 0
+
+    # First page
+    with requests_mock.Mocker() as mocker:
+        mocker.register_uri(
+            "GET",
+            "https://api.flare.io/leaksdb/v2/sources",
+            json={
+                "items": [1, 2, 3],
+                "next": "second_page",
+            },
+            status_code=200,
+        )
+        resp: requests.Response = next(response_iterator)
+        assert resp.json() == {"items": [1, 2, 3], "next": "second_page"}
+
+        assert len(mocker.request_history) == 1
+        assert mocker.last_request.query == ""
+
+    # Second Page
+    with requests_mock.Mocker() as mocker:
+        mocker.register_uri(
+            "GET",
+            "https://api.flare.io/leaksdb/v2/sources?from=second_page",
+            json={
+                "items": [4, 5, 6],
+                "next": "third_page",
+            },
+            status_code=200,
+        )
+        resp = next(response_iterator)
+        assert resp.json() == {"items": [4, 5, 6], "next": "third_page"}
+        assert len(mocker.request_history) == 1
+        assert mocker.last_request.query == "from=second_page"
+
+    # Third Page
+    with requests_mock.Mocker() as mocker:
+        mocker.register_uri(
+            "GET",
+            "https://api.flare.io/leaksdb/v2/sources?from=third_page",
+            json={
+                "items": [7, 8, 9],
+                "next": None,
+            },
+            status_code=200,
+        )
+        resp = next(response_iterator)
+        assert resp.json() == {
+            "items": [7, 8, 9],
+            "next": None,
+        }
+        assert len(mocker.request_history) == 1
+        assert mocker.last_request.query == "from=third_page"
+
+    # We stop here.
+    with requests_mock.Mocker() as mocker:
+        with pytest.raises(StopIteration):
+            resp = next(response_iterator)
+        assert len(mocker.request_history) == 0