fix(jumpstart): send NextToken when Hub.list_models pages through hub contents

src/sagemaker/jumpstart/hub/hub.py

@@ -126,6 +126,8 @@ def _list_and_paginate_models(self, **kwargs) -> List[Dict[str, Any]]:
 
         while first_iteration or next_token:
             first_iteration = False
+            if next_token:
+                kwargs["next_token"] = next_token
             list_hub_content_response = self._sagemaker_session.list_hub_contents(**kwargs)
             hub_model_summaries.extend(list_hub_content_response.get("HubContentSummaries", []))
             next_token = list_hub_content_response.get("NextToken")

tests/unit/sagemaker/jumpstart/hub/test_hub.py

@@ -212,6 +212,45 @@ def test_describe_model_one_thrown_error(mock_describe_hub_content_response, sag
         )
 
 
+def test_list_models_sends_next_token_to_the_following_page(sagemaker_session):
+    pages = {
+        ("ModelReference", None): {
+            "HubContentSummaries": [{"HubContentName": "reference-a"}],
+            "NextToken": "page-2",
+        },
+        ("ModelReference", "page-2"): {"HubContentSummaries": [{"HubContentName": "reference-b"}]},
+        ("Model", None): {
+            "HubContentSummaries": [{"HubContentName": "model-a"}],
+            "NextToken": "page-2",
+        },
+        ("Model", "page-2"): {
+            "HubContentSummaries": [{"HubContentName": "model-b"}],
+            "NextToken": "page-3",
+        },
+        ("Model", "page-3"): {"HubContentSummaries": [{"HubContentName": "model-c"}]},
+    }
+
+    def list_hub_contents(**kwargs):
+        assert kwargs["hub_name"] == HUB_NAME
+        assert kwargs["max_results"] == 2
+        # Each page is served once. A second request for a page raises KeyError, not a loop.
+        return pages.pop((kwargs["hub_content_type"], kwargs.get("next_token")))
+
+    sagemaker_session.list_hub_contents = Mock(side_effect=list_hub_contents)
+    hub = Hub(hub_name=HUB_NAME, sagemaker_session=sagemaker_session)
+
+    response = hub.list_models(max_results=2)
+
+    assert [summary["HubContentName"] for summary in response["hub_content_summaries"]] == [
+        "reference-a",
+        "reference-b",
+        "model-a",
+        "model-b",
+        "model-c",
+    ]
+    assert pages == {}
+
+
 def test_create_hub_content_reference(sagemaker_session):
     hub = Hub(hub_name=HUB_NAME, sagemaker_session=sagemaker_session)
     model_name = "mock-model-one-huggingface"