update session and add tests

Author: bencrabtree

src/sagemaker/session.py

@@ -6485,14 +6485,20 @@ def create_hub(
         Returns:
             (dict): Return value from the ``CreateHub`` API.
         """
-        return self.sagemaker_client.create_hub(
-            hub_name=hub_name,
-            hub_description=hub_description,
-            hub_display_name=hub_display_name,
-            hub_search_keywords=hub_search_keywords,
-            s3_storage_config=s3_storage_config,
-            tags=tags
-        )
+        request = {
+            "HubName": hub_name,
+            "HubDescription": hub_description
+        }
+        if hub_display_name:
+            request["HubDisplayName"] = hub_display_name
+        if hub_search_keywords:
+            request["HubSearchKeywords"] = hub_search_keywords
+        if s3_storage_config:
+            request["S3StorageConfig"] = s3_storage_config
+        if tags:
+            request["Tags"] = tags
+
+        return self.sagemaker_client.create_hub(**request)
 
     def describe_hub(
         self,
@@ -6506,9 +6512,9 @@ def describe_hub(
         Returns:
             (dict): Return value for ``DescribeHub`` API
         """
-        return self.sagemaker_client.describe_hub(
-            hub_name=hub_name
-        )
+        request = { "HubName": hub_name }
+
+        return self.sagemaker_client.describe_hub(**request)
 
     def list_hubs(
         self,
@@ -6535,15 +6541,23 @@ def list_hubs(
         Returns:
             (dict): Return value for ``ListHubs`` API
         """
-        return self.sagemaker_client.list_hubs(
-            creation_time_after=creation_time_after,
-            creation_time_before=creation_time_before,
-            max_results=max_results,
-            max_schema_version=max_schema_version,
-            name_contains=name_contains,
-            sort_by=sort_by,
-            sort_order=sort_order
-        )
+        request = {}
+        if creation_time_after:
+            request["CreationTimeAfter"] = creation_time_after
+        if creation_time_before:
+            request["CreationTimeBefore"] = creation_time_before
+        if max_results:
+            request["MaxResults"] = max_results
+        if max_schema_version:
+            request["MaxSchemaVersion"] = max_schema_version
+        if name_contains:
+            request["NameContains"] = name_contains
+        if sort_by:
+            request["SortBy"] = sort_by
+        if sort_order:
+            request["SortOrder"] = sort_order
+
+        return self.sagemaker_client.list_hubs(*request)
 
     def list_hub_contents(
         self,
@@ -6574,17 +6588,26 @@ def list_hub_contents(
         Returns:
             (dict): Return value for ``ListHubContents`` API
         """
-        return self.sagemaker_client.list_hub_contents(
-            hub_name=hub_name,
-            hub_content_type=hub_content_type,
-            creation_time_after=creation_time_after,
-            creation_time_before=creation_time_before,
-            max_results=max_results,
-            max_schema_version=max_schema_version,
-            name_contains=name_contains,
-            sort_by=sort_by,
-            sort_order=sort_order
-        )
+        request = {
+            "HubName": hub_name,
+            "HubContentType": hub_content_type
+        }
+        if creation_time_after:
+            request["CreationTimeAfter"] = creation_time_after
+        if creation_time_before:
+            request["CreationTimeBefore"] = creation_time_before
+        if max_results:
+            request["MaxResults"] = max_results
+        if max_schema_version:
+            request["MaxSchemaVersion"] = max_schema_version
+        if name_contains:
+            request["NameContains"] = name_contains
+        if sort_by:
+            request["SortBy"] = sort_by
+        if sort_order:
+            request["SortOrder"] = sort_order
+
+        return self.sagemaker_client.list_hub_contents(*request)
 
     def delete_hub(
         self,
@@ -6595,14 +6618,17 @@ def delete_hub(
         Args:
             hub_name (str): The name of the hub to delete.
         """
-        return self.sagemaker_client.delete_hub(hub_name=hub_name)
+        request = { "HubName": hub_name }
+
+        return self.sagemaker_client.delete_hub(*request)
 
     def import_hub_content(
         self,
         document_schema_version: str,
         hub_content_name: str,
         hub_content_type: str,
         hub_name: str,
+        hub_content_document: str,
         hub_content_display_name=None,
         hub_content_description=None,
         hub_content_version=None,
@@ -6618,6 +6644,8 @@ def import_hub_content(
             hub_content_version (str): The version of the HubContent to import.
             hub_content_type (str): The type of HubContent to import.
             hub_name (str): The name of the Hub to import content to.
+            hub_content_document (str): The hub content document that describes information about the hub content
+                such as type, associated containers, scripts, and more.
             hub_content_display_name (str): The display name of the HubContent to import.
             hub_content_description (str): The description of the HubContent to import.
             hub_content_markdown (str): A string that provides a description of the HubContent. This string can include links, tables,
@@ -6627,18 +6655,27 @@ def import_hub_content(
         Returns:
             (dict): Return value for ``ImportHubContent`` API
         """
-        return self.sagemaker_client.import_hub_content(
-            document_schema_version=document_schema_version,
-            hub_content_name=hub_content_name,
-            hub_content_version=hub_content_version,
-            hub_content_type=hub_content_type,
-            hub_name=hub_name,
-            hub_content_display_name=hub_content_display_name,
-            hub_content_description=hub_content_description,
-            hub_content_markdown=hub_content_markdown,
-            hub_content_search_keywords=hub_content_search_keywords,
-            tags=tags
-        )
+        request = {
+            "DocumentSchemaVersion": document_schema_version,
+            "HubContentName": hub_content_name,
+            "HubContentType": hub_content_type,
+            "HubName": hub_name,
+            "HubContentDocument": hub_content_document
+        }
+        if hub_content_display_name:
+            request["HubContentDisplayName"] = hub_content_display_name
+        if hub_content_description:
+            request["HubContentDescription"] = hub_content_description
+        if hub_content_version:
+            request["HubContentVersion"] = hub_content_version
+        if hub_content_markdown:
+            request["HubContentMarkdown"] = hub_content_markdown
+        if hub_content_search_keywords:
+            request["HubContentSearchKeywords"] = hub_content_search_keywords
+        if tags:
+            request["Tags"] = tags
+        
+        return self.sagemaker_client.import_hub_content(*request)
 
     def describe_hub_content(
         self,
@@ -6658,12 +6695,15 @@ def describe_hub_content(
         Returns:
             (dict): Return value for ``DescribeHubContent`` API
         """
-        return self.sagemaker_client.describe_hub_content(
-            hub_content_name=hub_content_name,
-            hub_content_type=hub_content_type,
-            hub_name=hub_name,
-            hub_content_version=hub_content_version
-        )
+        request = {
+            "HubContentName": hub_content_name,
+            "HubContentType": hub_content_type,
+            "HubName": hub_name
+        }
+        if hub_content_version:
+            request["HubContentVersion"] = hub_content_version
+
+        return self.sagemaker_client.describe_hub_content(*request)
 
     def delete_hub_content(
         self,
@@ -6680,12 +6720,14 @@ def delete_hub_content(
             hub_content_type (str): The type of the content that you want to delete from a Hub.
             hub_name (str): The name of the Hub that you want to delete content in.
         """
-        return self.sagemaker_client.delete_hub_content(
-            hub_content_name=hub_content_name,
-            hub_content_version=hub_content_version,
-            hub_content_type=hub_content_type,
-            hub_name=hub_name
-        )
+        request = {
+            "HubContentName": hub_content_name,
+            "HubContentType": hub_content_type,
+            "HubName": hub_name,
+            "HubContentVersion": hub_content_version
+        }
+
+        return self.sagemaker_client.delete_hub_content(*request)
 
 
 def get_model_package_args(

tests/unit/test_session.py

@@ -6489,3 +6489,163 @@ def test_download_data_with_file_and_directory(makedirs, sagemaker_session):
         Filename="./foo/bar/mode.tar.gz",
         ExtraArgs=None,
     )
+
+def test_create_hub(sagemaker_session):
+    sagemaker_session.create_hub(
+        hub_name="mock-hub-name",
+        hub_description="this is my sagemaker hub",
+        hub_display_name="Mock Hub",
+        hub_search_keywords=["mock", "hub", "123"],
+        s3_storage_config={
+            "S3OutputPath": "s3://my-hub-bucket/"
+        },
+        tags=[{"Key": "tag-key-1", "Value": "tag-value-1"}]
+    )
+
+    request = {
+        "hub_name": "mock-hub-name",
+        "hub_description": "this is my sagemaker hub",
+        "hub_display_name": "Mock Hub",
+        "hub_search_keywords": ["mock", "hub", "123"],
+        "s3_storage_config": {
+            "S3OutputPath": "s3://my-hub-bucket/"
+        },
+        "tags": [{"Key": "tag-key-1", "Value": "tag-value-1"}]
+    }
+
+    sagemaker_session.create_hub.assert_called_with(**request)
+
+def test_describe_hub(sagemaker_session):
+    sagemaker_session.describe_hub(
+        hub_name="mock-hub-name",
+    )
+
+    request = {
+        "hub_name": "mock-hub-name",
+    }
+
+    sagemaker_session.describe_hub.assert_called_with(**request)
+
+def test_list_hubs(sagemaker_session):
+    sagemaker_session.list_hubs(
+        creation_time_after="08-14-1997 12:00:00",
+        creation_time_before="01-08/2024 10:25:00",
+        max_results="25",
+        max_schema_version="1.0.5",
+        name_contains="mock-hub",
+        sort_by="HubName",
+        sort_order="Ascending"
+    )
+
+    request = {
+        "creation_time_after": "08-14-1997 12:00:00",
+        "creation_time_before": "01-08/2024 10:25:00",
+        "max_results": "25",
+        "max_schema_version": "1.0.5",
+        "name_contains": "mock-hub",
+        "sort_by": "HubName",
+        "sort_order": "Ascending"
+    }
+
+    sagemaker_session.list_hubs.assert_called_with(**request)
+
+def test_list_hub_contents(sagemaker_session):
+    sagemaker_session.list_hub_contents(
+        hub_name="mock-hub-123",
+        hub_content_type="MODEL",
+        creation_time_after="08-14-1997 12:00:00",
+        creation_time_before="01-08/2024 10:25:00",
+        max_results="25",
+        max_schema_version="1.0.5",
+        name_contains="mock-hub",
+        sort_by="HubName",
+        sort_order="Ascending"
+    )
+
+    request = {
+        "hub_name": "mock-hub-123",
+        "hub_content_type": "MODEL",
+        "creation_time_after": "08-14-1997 12:00:00",
+        "creation_time_before": "01-08/2024 10:25:00",
+        "max_results": "25",
+        "max_schema_version": "1.0.5",
+        "name_contains": "mock-hub",
+        "sort_by": "HubName",
+        "sort_order": "Ascending"
+    }
+
+    sagemaker_session.list_hub_contents.assert_called_with(**request)
+
+def test_delete_hub(sagemaker_session):
+    sagemaker_session.delete_hub(
+        hub_name="mock-hub-123",
+    )
+
+    request = {
+        "hub_name": "mock-hub-123",
+    }
+
+    sagemaker_session.delete_hub.assert_called_with(**request)
+
+def test_import_hub_content(sagemaker_session):
+    sagemaker_session.import_hub_content(
+        hub_name="mock-hub-123",
+        hub_content_type="MODEL",
+        document_schema_version="1.0.0",
+        hub_content_document="{'training_script_location':'s3://path/to/script.py'}",
+        hub_content_name="mock-hub-content-1",
+        hub_content_display_name="Mock Hub Content One",
+        hub_content_description="This is my special Hub Content for my special Hub",
+        hub_content_version="5.5.5",
+        hub_content_markdown="markdown",
+        hub_content_search_keywords=["Hub","Machine Learning","Content"]
+    )
+
+    request = {
+        "hub_name": "mock-hub-123",
+        "hub_content_type": "MODEL",
+        "document_schema_version": "1.0.0",
+        "hub_content_document": "{'training_script_location':'s3://path/to/script.py'}",
+        "hub_content_name": "mock-hub-content-1",
+        "hub_content_display_name": "Mock Hub Content One",
+        "hub_content_description": "This is my special Hub Content for my special Hub",
+        "hub_content_version": "5.5.5",
+        "hub_content_markdown": "markdown",
+        "hub_content_search_keywords": ["Hub","Machine Learning","Content"]
+    }
+
+    sagemaker_session.import_hub_content.assert_called_with(**request)
+
+def test_describe_hub_content(sagemaker_session):
+    sagemaker_session.describe_hub_content(
+        hub_name="mock-hub-123",
+        hub_content_type="MODEL",
+        hub_content_name="mock-hub-content-1",
+        hub_content_version="5.5.5",
+    )
+
+    request = {
+        "hub_name": "mock-hub-123",
+        "hub_content_type": "MODEL",
+        "hub_content_name": "mock-hub-content-1",
+        "hub_content_version": "5.5.5",
+    }
+
+    sagemaker_session.describe_hub_content.assert_called_with(**request)
+
+def test_delete_hub_content(sagemaker_session):
+    sagemaker_session.delete_hub_content(
+        hub_name="mock-hub-123",
+        hub_content_type="MODEL",
+        hub_content_name="mock-hub-content-1",
+        hub_content_version="5.5.5",
+    )
+
+    request = {
+        "hub_name": "mock-hub-123",
+        "hub_content_type": "MODEL",
+        "hub_content_name": "mock-hub-content-1",
+        "hub_content_version": "5.5.5",
+    }
+
+    sagemaker_session.delete_hub_content.assert_called_with(**request)