test_tag/test_tags method assert fix in association tests

Author: stisac

tests/integ/sagemaker/lineage/test_association.py

@@ -66,7 +66,9 @@ def test_set_tag(association_obj, sagemaker_session):
         if actual_tags:
             break
         time.sleep(1)
-    assert len(actual_tags) == 1
+    # When sagemaker-client-config endpoint-url is passed as argument to hit beta,
+    # length of actual tags will be 2
+    assert len(actual_tags) > 0
     assert actual_tags[0] == tag
 
 
@@ -81,5 +83,7 @@ def test_tags(association_obj, sagemaker_session):
         if actual_tags:
             break
         time.sleep(1)
-    assert len(actual_tags) == 1
-    assert actual_tags == tags
+    # When sagemaker-client-config endpoint-url is passed as argument to hit beta,
+    # length of actual tags will be 2
+    assert len(actual_tags) > 0
+    assert [actual_tags[-1]] == tags