Skip to content

Commit 1401e62

Browse files
committed
add integration test for listing artifacts by type
1 parent 131947f commit 1401e62

File tree

2 files changed

+30
-0
lines changed

2 files changed

+30
-0
lines changed

tests/integ/sagemaker/lineage/conftest.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,16 @@ def artifact_objs(sagemaker_session):
302302
)
303303
time.sleep(SLEEP_TIME_SECONDS)
304304

305+
artifact_objs.append(
306+
artifact.Artifact.create(
307+
artifact_name=name(),
308+
artifact_type="SDKIntegrationTestType2",
309+
source_uri=name(),
310+
properties={"k1": "v1"},
311+
sagemaker_session=sagemaker_session,
312+
)
313+
)
314+
305315
yield artifact_objs
306316

307317
for artifact_obj in artifact_objs:

tests/integ/sagemaker/lineage/test_artifact.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,26 @@ def test_list(artifact_objs, sagemaker_session):
7777
assert artifact_names
7878

7979

80+
def test_list_by_type(artifact_objs, sagemaker_session):
81+
slack = datetime.timedelta(minutes=1)
82+
now = datetime.datetime.now(datetime.timezone.utc)
83+
expected_name = list(filter(lambda x: x.artifact_type == 'SDKIntegrationTestType2', artifact_objs))[0].artifact_name
84+
artifact_names = [art.artifact_name for art in artifact_objs]
85+
86+
artifact_names_listed = [
87+
artifact_listed.artifact_name
88+
for artifact_listed in artifact.Artifact.list(
89+
created_after=now - slack,
90+
artifact_type='SDKIntegrationTestType2',
91+
sagemaker_session=sagemaker_session,
92+
)
93+
if artifact_listed.artifact_name in artifact_names
94+
]
95+
96+
assert len(artifact_names_listed) == 1
97+
assert artifact_names_listed[0] == expected_name
98+
99+
80100
def test_downstream_trials(trial_associated_artifact, trial_obj, sagemaker_session):
81101
# allow trial components to index, 30 seconds max
82102
for i in range(3):

0 commit comments

Comments
 (0)