From b8019965340938f500ff51ceaf4cdf8a7cb253b8 Mon Sep 17 00:00:00 2001
From: Dana Benson <danabens@amazon.com>
Date: Thu, 28 Jan 2021 09:45:14 -0800
Subject: [PATCH] add integration test for listing artifacts by type

---
 tests/integ/sagemaker/lineage/conftest.py     | 10 +++++++++
 .../integ/sagemaker/lineage/test_artifact.py  | 22 +++++++++++++++++++
 2 files changed, 32 insertions(+)

diff --git a/tests/integ/sagemaker/lineage/conftest.py b/tests/integ/sagemaker/lineage/conftest.py
index 6671fb6a99..9fd0fa2f73 100644
--- a/tests/integ/sagemaker/lineage/conftest.py
+++ b/tests/integ/sagemaker/lineage/conftest.py
@@ -302,6 +302,16 @@ def artifact_objs(sagemaker_session):
         )
         time.sleep(SLEEP_TIME_SECONDS)
 
+    artifact_objs.append(
+        artifact.Artifact.create(
+            artifact_name=name(),
+            artifact_type="SDKIntegrationTestType2",
+            source_uri=name(),
+            properties={"k1": "v1"},
+            sagemaker_session=sagemaker_session,
+        )
+    )
+
     yield artifact_objs
 
     for artifact_obj in artifact_objs:
diff --git a/tests/integ/sagemaker/lineage/test_artifact.py b/tests/integ/sagemaker/lineage/test_artifact.py
index ffdfc9dc81..84e32e2575 100644
--- a/tests/integ/sagemaker/lineage/test_artifact.py
+++ b/tests/integ/sagemaker/lineage/test_artifact.py
@@ -77,6 +77,28 @@ def test_list(artifact_objs, sagemaker_session):
     assert artifact_names
 
 
+def test_list_by_type(artifact_objs, sagemaker_session):
+    slack = datetime.timedelta(minutes=1)
+    now = datetime.datetime.now(datetime.timezone.utc)
+    expected_name = list(
+        filter(lambda x: x.artifact_type == "SDKIntegrationTestType2", artifact_objs)
+    )[0].artifact_name
+    artifact_names = [art.artifact_name for art in artifact_objs]
+
+    artifact_names_listed = [
+        artifact_listed.artifact_name
+        for artifact_listed in artifact.Artifact.list(
+            created_after=now - slack,
+            artifact_type="SDKIntegrationTestType2",
+            sagemaker_session=sagemaker_session,
+        )
+        if artifact_listed.artifact_name in artifact_names
+    ]
+
+    assert len(artifact_names_listed) == 1
+    assert artifact_names_listed[0] == expected_name
+
+
 def test_downstream_trials(trial_associated_artifact, trial_obj, sagemaker_session):
     # allow trial components to index, 30 seconds max
     for i in range(3):