Limit page nexts (#4)

Author: jaschultz

tests/test_testspace_api.py

@@ -86,12 +86,42 @@ def test_get_projects_paginated(load_json, testspace_api, requests_mock):
         json=load_json[:1],
         headers={"link": ", ".join([links_string_first, links_string_last])},
     )
-    response_json = testspace_api.get_projects()
+    response_json = testspace_api.get_projects(limit=None)
 
     project_names = [project["name"] for project in response_json]
     assert testspace_api.project in project_names
 
 
+@pytest.mark.parametrize("load_json", ["projects.json"], indirect=True)
+def test_get_projects_paginated_limited(load_json, testspace_api, requests_mock):
+    url = testspace_api.url
+    testspace_url = "https://{}/api/projects".format(url)
+
+    links_string_first = '<{}?page={}>; rel="{}"'.format(testspace_url, 1, "first")
+    links_string_next = '<{}?page={}>; rel="{}"'.format(testspace_url, 2, "next")
+    links_string_last = '<{}?page={}>; rel="{}"'.format(testspace_url, 2, "last")
+
+    requests_mock.get(
+        "/api/projects",
+        json=load_json[1:],
+        headers={
+            "link": ", ".join(
+                [links_string_first, links_string_next, links_string_last]
+            )
+        },
+        complete_qs=True,
+    )
+    requests_mock.get(
+        "/api/projects?page=2",
+        json=load_json[:1],
+        headers={"link": ", ".join([links_string_first, links_string_last])},
+    )
+    num_projects = 1
+    response_json = testspace_api.get_projects(limit=num_projects)
+
+    assert len(response_json) == num_projects
+
+
 @pytest.mark.parametrize("load_json", ["projects.json"], indirect=True)
 def test_get_project_name(load_json, testspace_api, requests_mock):
     project = testspace_api.project

testspace/testspace.py

@@ -174,19 +174,22 @@ def get_request(self, path=None):
         return response
 
     def get_request_json(self, path=None, limit=30):
-        if type(limit) is not int or limit <= 0:
+        if limit is None:
+            pass
+        elif type(limit) is not int or limit <= 0:
             raise ValueError
         response = self.get_request(path)
-        if len(response.links) == 0:
-            return response.json()
-        else:
+        response_json = response.json()
+        if type(response_json) is list:
             next_url = response.links.get("next", None)
-            response_json = response.json()
             while next_url is not None:
                 response = self.get_request(next_url.get("url"))
                 next_url = response.links.get("next", None)
                 response_json.extend(response.json())
-        return response_json[:limit]
+                if limit is not None and len(response_json) >= limit:
+                    break
+            response_json = response_json[:limit]
+        return response_json
 
     def post_request(self, payload, path=None):
         response = self._api_request("POST", path=path, payload=payload)