Skip to content

Commit 50fa48e

Browse files
committed
Add unit test for construct hub model reference arn util
1 parent 280efe0 commit 50fa48e

File tree

1 file changed

+14
-0
lines changed

1 file changed

+14
-0
lines changed

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

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -95,6 +95,20 @@ def test_construct_hub_model_arn_from_inputs():
9595
== "arn:aws:sagemaker:us-west-2:123456789123:hub-content/my-mock-hub/Model/pytorch-ic-imagenet-v2/*"
9696
)
9797

98+
def test_construct_hub_model_reference_arn_from_inputs():
99+
model_name, version = "pytorch-ic-imagenet-v2", "1.0.2"
100+
hub_arn = "arn:aws:sagemaker:us-west-2:123456789123:hub/my-mock-hub"
101+
102+
assert (
103+
utils.construct_hub_model_arn_from_inputs(hub_arn, model_name, version)
104+
== "arn:aws:sagemaker:us-west-2:123456789123:hub-content/my-mock-hub/ModelReference/pytorch-ic-imagenet-v2/1.0.2"
105+
)
106+
107+
version = "*"
108+
assert (
109+
utils.construct_hub_model_arn_from_inputs(hub_arn, model_name, version)
110+
== "arn:aws:sagemaker:us-west-2:123456789123:hub-content/my-mock-hub/ModelReference/pytorch-ic-imagenet-v2/*"
111+
)
98112

99113
def test_generate_hub_arn_for_init_kwargs():
100114
hub_name = "my-hub-name"

0 commit comments

Comments
 (0)