Remove pytest fixture and fix test_tag/s method

tests/integ/sagemaker/lineage/conftest.py

@@ -25,19 +25,14 @@
     association,
     artifact,
 )
-from sagemaker.session import Session
+
 from smexperiments import trial_component, trial, experiment
 
 from tests.integ.sagemaker.lineage.helpers import name, names
 
 SLEEP_TIME_SECONDS = 1
 
 
-@pytest.fixture()
-def sagemaker_session(boto_session):
-    return Session(boto_session=boto_session)
-
-
 @pytest.fixture
 def action_obj(sagemaker_session):
     obj = action.Action.create(

tests/integ/sagemaker/lineage/test_action.py

@@ -90,7 +90,9 @@ def test_tag(action_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-    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
 
 
@@ -104,5 +106,7 @@ def test_tags(action_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-    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

tests/integ/sagemaker/lineage/test_artifact.py

@@ -99,7 +99,9 @@ def test_tag(artifact_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-    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
 
 
@@ -113,5 +115,7 @@ def test_tags(artifact_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-    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

tests/integ/sagemaker/lineage/test_context.py

@@ -88,7 +88,9 @@ def test_tag(context_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-    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
 
 
@@ -102,5 +104,7 @@ def test_tags(context_obj, sagemaker_session):
         )["Tags"]
         if actual_tags:
             break
-    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