Fix List Associations integ tests

Author: stisac

tests/integ/sagemaker/lineage/test_association.py

@@ -15,17 +15,14 @@
 
 import datetime
 
-import pytest
 from sagemaker.lineage import association
 
 
-@pytest.mark.skip(reason="Not in CMH yet")
 def test_create_delete(association_obj):
     # fixture does create and then delete, this test ensures it happens at least once
     assert association_obj.source_arn
 
 
-@pytest.mark.skip(reason="Not in CMH yet")
 def test_list(association_objs, sagemaker_session):
     slack = datetime.timedelta(minutes=1)
     now = datetime.datetime.now(datetime.timezone.utc)
@@ -35,27 +32,51 @@ def test_list(association_objs, sagemaker_session):
 
     for sort_order in ["Ascending", "Descending"]:
         association_keys_listed = []
+        source_arn = [assoc_obj.source_arn for assoc_obj in association_objs][0]
         listed = association.Association.list(
+            source_arn=source_arn,
             created_after=now - slack,
             created_before=now + slack,
             sort_by="CreationTime",
             sort_order=sort_order,
             sagemaker_session=sagemaker_session,
         )
+
         for assoc in listed:
             key = assoc.source_arn + ":" + assoc.destination_arn
             if key in association_keys:
                 association_keys_listed.append(key)
 
     if sort_order == "Descending":
         association_names_listed = association_keys_listed[::-1]
-    assert association_keys == association_names_listed
+    assert association_keys[::-1] == association_names_listed
     # sanity check
     assert association_keys_listed
 
 
-@pytest.mark.skip(reason="Not in CMH yet")
 def test_set_tag(association_obj, sagemaker_session):
     tag = {"Key": "foo", "Value": "bar"}
     association_obj.set_tag(tag)
-    assert association_obj.get_tag() == tag
+
+    while True:
+        actual_tags = sagemaker_session.sagemaker_client.list_tags(
+            ResourceArn=association_obj.source_arn
+        )["Tags"]
+        if actual_tags:
+            break
+    assert len(actual_tags) == 1
+    assert actual_tags[0] == tag
+
+
+def test_tags(association_obj, sagemaker_session):
+    tags = [{"Key": "foo1", "Value": "bar1"}]
+    association_obj.set_tags(tags)
+
+    while True:
+        actual_tags = sagemaker_session.sagemaker_client.list_tags(
+            ResourceArn=association_obj.source_arn
+        )["Tags"]
+        if actual_tags:
+            break
+    assert len(actual_tags) == 1
+    assert actual_tags == tags