File tree 1 file changed +3
-0
lines changed
tests/unit/sagemaker/jumpstart/hub
1 file changed +3
-0
lines changed Original file line number Diff line number Diff line change 17
17
from sagemaker .jumpstart .hub .interfaces import HubModelDocument
18
18
from tests .unit .sagemaker .jumpstart .constants import HUB_MODEL_DOCUMENT_DICTS
19
19
from unittest .mock import MagicMock
20
+ from sagemaker .jumpstart .types import HubContentType
21
+
20
22
21
23
REGION = "us-east-1"
22
24
ACCOUNT_ID = "123456789123"
@@ -41,6 +43,7 @@ def test_parse_(input_string, expected):
41
43
def test_make_model_specs_from_describe_hub_content_response ():
42
44
mock_describe_response = MagicMock ()
43
45
region = "us-west-2"
46
+ mock_describe_response .hub_content_type = HubContentType .MODEL
44
47
mock_describe_response .get_hub_region .return_value = region
45
48
mock_describe_response .hub_content_version = "1.0.0"
46
49
json_obj = HUB_MODEL_DOCUMENT_DICTS ["huggingface-llm-gemma-2b-instruct" ]
You can’t perform that action at this time.
0 commit comments