diff --git a/runpod/api/ctl_commands.py b/runpod/api/ctl_commands.py index 1291df490..4182d9129 100644 --- a/runpod/api/ctl_commands.py +++ b/runpod/api/ctl_commands.py @@ -18,6 +18,20 @@ def _path_segment(value: str) -> str: return quote(value, safe="") +def _list_all( + path: str, collection: str, *, api_key: Optional[str] = None +) -> list[dict]: + records = [] + params = {} + while True: + response = run_rest_request("GET", path, api_key=api_key, params=params) + records.extend(response[collection]) + pagination = response["pagination"] + if not pagination["hasNextPage"]: + return records + params["cursor"] = pagination["nextCursor"] + + def _split_values(value: Optional[Iterable[Any] | str]) -> list[str]: if value is None: return [] @@ -93,8 +107,7 @@ def get_gpu(gpu_id: str, gpu_quantity: int = 1, api_key: Optional[str] = None) - def get_pods(api_key: Optional[str] = None) -> list[dict]: """Get all standalone pods.""" - response = run_rest_request("GET", "/v2/pods", api_key=api_key) - return response["pods"] + return _list_all("/v2/pods", "pods", api_key=api_key) def get_pod(pod_id: str, api_key: Optional[str] = None) -> Optional[dict]: @@ -279,8 +292,7 @@ def create_template( def get_endpoints() -> list[dict]: """Get all serverless endpoints.""" - response = run_rest_request("GET", "/v2/serverless") - return response["endpoints"] + return _list_all("/v2/serverless", "endpoints") def create_endpoint( diff --git a/tests/test_api/test_ctl_commands.py b/tests/test_api/test_ctl_commands.py index d3134abc8..13185fb96 100644 --- a/tests/test_api/test_ctl_commands.py +++ b/tests/test_api/test_ctl_commands.py @@ -90,11 +90,12 @@ def test_get_gpu_propagates_other_api_errors(): def test_get_pods_unwraps_response(): pods = [{"id": "pod"}] with patch( - "runpod.api.ctl_commands.run_rest_request", return_value={"pods": pods} + "runpod.api.ctl_commands.run_rest_request", + return_value={"pods": pods, "pagination": {"hasNextPage": False}}, ) as request: assert ctl_commands.get_pods(api_key="key") == pods - request.assert_called_once_with("GET", "/v2/pods", api_key="key") + request.assert_called_once_with("GET", "/v2/pods", api_key="key", params={}) def test_get_pod_escapes_id(): @@ -343,11 +344,11 @@ def test_get_endpoints_unwraps_response(): endpoints = [{"id": "endpoint"}] with patch( "runpod.api.ctl_commands.run_rest_request", - return_value={"endpoints": endpoints}, + return_value={"endpoints": endpoints, "pagination": {"hasNextPage": False}}, ) as request: assert ctl_commands.get_endpoints() == endpoints - request.assert_called_once_with("GET", "/v2/serverless") + request.assert_called_once_with("GET", "/v2/serverless", api_key=None, params={}) @pytest.fixture diff --git a/tests/test_api/test_rest.py b/tests/test_api/test_rest.py index 5e133aedf..673cef993 100644 --- a/tests/test_api/test_rest.py +++ b/tests/test_api/test_rest.py @@ -27,6 +27,7 @@ def rest_server(monkeypatch): "status": 200, "body": b'{"id":"pod-1"}', "content_type": "application/json", + "requests": [], "started": threading.Event(), "disconnected": threading.Event(), } @@ -34,6 +35,7 @@ def rest_server(monkeypatch): class Handler(BaseHTTPRequestHandler): def _respond(self): state["hits"] += 1 + state["requests"].append(self.path) body = self.rfile.read(int(self.headers.get("Content-Length", 0))) state["started"].set() if state.get("stall"): @@ -48,6 +50,10 @@ def _respond(self): status, response_body = state["status"], state["body"] if self.headers.get("Authorization") != "Bearer key": status, response_body = 401, b"invalid key" + elif "responses" in state: + status, response_body = state["responses"].get( + self.path, (404, b"wrong path") + ) elif self.path != state.get("path", "/prefix/v2/pods"): status, response_body = 404, b"wrong path" elif ( @@ -109,6 +115,49 @@ def test_authenticated_mutation_at_custom_base_url( assert result == {"id": "pod-1"} +@pytest.mark.parametrize( + ("get_records", "path", "collection", "kwargs"), + [ + (runpod.get_pods, "/v2/pods", "pods", {"api_key": "key"}), + (runpod.get_pods, "/v2/pods", "pods", {}), + (runpod.get_endpoints, "/v2/serverless", "endpoints", {}), + ], +) +@pytest.mark.parametrize("second_page_status", [200, 503]) +def test_sdk_lists_follow_cursor_pages( + rest_server, monkeypatch, get_records, path, collection, kwargs, second_page_status +): + monkeypatch.setattr(runpod, "api_key", "wrong-key" if kwargs else "key") + first_path = f"/prefix{path}" + second_path = f"{first_path}?cursor=next-page" + second_body = ( + {collection: [{"id": "second"}], "pagination": {"hasNextPage": False}} + if second_page_status == 200 + else {"detail": "second page unavailable"} + ) + rest_server["responses"] = { + first_path: ( + 200, + json.dumps( + { + collection: [{"id": "first"}], + "pagination": {"hasNextPage": True, "nextCursor": "next-page"}, + } + ).encode(), + ), + second_path: (second_page_status, json.dumps(second_body).encode()), + } + + if second_page_status == 200: + assert get_records(**kwargs) == [{"id": "first"}, {"id": "second"}] + else: + with pytest.raises(QueryError, match="second page unavailable") as raised: + get_records(**kwargs) + assert raised.value.status_code == 503 + + assert rest_server["requests"] == [first_path, second_path] + + def test_missing_key_fails_before_network_request(rest_server, rest_request): with pytest.raises(AuthenticationError): rest_request("POST", "/v2/pods", json={"name": "pod"})