Skip to content

Commit 11c5160

Browse files
committed
Add unit test
1 parent f6d9538 commit 11c5160

File tree

1 file changed

+30
-0
lines changed

1 file changed

+30
-0
lines changed

tests/unit/sagemaker/model/test_model.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1307,3 +1307,33 @@ def test_package_for_edge_with_sagemaker_config_injection(sagemaker_session):
13071307
role=SAGEMAKER_CONFIG_EDGE_PACKAGING_JOB["SageMaker"]["EdgePackagingJob"]["RoleArn"],
13081308
tags=None,
13091309
)
1310+
1311+
1312+
def test_model_source(
1313+
sagemaker_session,
1314+
):
1315+
model = Model(
1316+
entry_point=ENTRY_POINT_INFERENCE,
1317+
role=ROLE,
1318+
sagemaker_session=sagemaker_session,
1319+
image_uri=IMAGE_URI,
1320+
model_data={
1321+
"S3DataSource": {
1322+
"S3Uri": "s3://tmybuckaet",
1323+
"S3DataType": "S3Prefix",
1324+
"CompressionType": "None",
1325+
}
1326+
},
1327+
)
1328+
1329+
assert model._get_model_uri() == "s3://tmybuckaet"
1330+
1331+
model_1 = Model(
1332+
entry_point=ENTRY_POINT_INFERENCE,
1333+
role=ROLE,
1334+
sagemaker_session=sagemaker_session,
1335+
image_uri=IMAGE_URI,
1336+
model_data="s3://tmybuckaet",
1337+
)
1338+
1339+
assert model_1._get_model_uri() == "s3://tmybuckaet"

0 commit comments

Comments
 (0)