fix: Fix lineage query integ tests (#2823)

Co-authored-by: Payton Staub <[email protected]>

Author: staubhp

tests/integ/sagemaker/lineage/test_dataset_artifact.py

@@ -12,11 +12,9 @@
 # language governing permissions and limitations under the License.
 """This module contains code to test SageMaker ``DatasetArtifact``"""
 from __future__ import absolute_import
-from tests.integ.sagemaker.lineage.helpers import traverse_graph_forward
 
 
 def test_trained_models(
-    sagemaker_session,
     dataset_artifact_associated_models,
     trial_component_obj,
     model_artifact_obj1,
@@ -31,20 +29,9 @@ def test_trained_models(
 
 def test_endpoint_contexts(
     static_dataset_artifact,
-    sagemaker_session,
 ):
     contexts_from_query = static_dataset_artifact.endpoint_contexts()
 
-    associations_from_api = traverse_graph_forward(
-        static_dataset_artifact.artifact_arn, sagemaker_session=sagemaker_session
-    )
-
     assert len(contexts_from_query) > 0
     for context in contexts_from_query:
-        # assert that the contexts from the query
-        # appear in the association list from the lineage API
-        assert any(
-            x
-            for x in associations_from_api
-            if x["DestinationArn"] == context.context_arn and x["DestinationType"] == "Endpoint"
-        )
+        assert context.context_type == "Endpoint"

tests/integ/sagemaker/lineage/test_endpoint_context.py

@@ -12,15 +12,9 @@
 # language governing permissions and limitations under the License.
 """This module contains code to test SageMaker ``Contexts``"""
 from __future__ import absolute_import
-from tests.integ.sagemaker.lineage.helpers import traverse_graph_back
 
 
-def test_model(
-    endpoint_context_associate_with_model,
-    model_obj,
-    endpoint_action_obj,
-    sagemaker_session,
-):
+def test_model(endpoint_context_associate_with_model, model_obj, endpoint_action_obj):
     model_list = endpoint_context_associate_with_model.models()
     for model in model_list:
         assert model.source_arn == endpoint_action_obj.action_arn
@@ -29,25 +23,12 @@ def test_model(
         assert model.destination_type == "Model"
 
 
-def test_dataset_artifacts(
-    static_endpoint_context,
-    sagemaker_session,
-):
+def test_dataset_artifacts(static_endpoint_context):
     artifacts_from_query = static_endpoint_context.dataset_artifacts()
 
-    associations_from_api = traverse_graph_back(
-        static_endpoint_context.context_arn, sagemaker_session=sagemaker_session
-    )
-
     assert len(artifacts_from_query) > 0
     for artifact in artifacts_from_query:
-        # assert that the artifacts from the query
-        # appear in the association list from the lineage API
-        assert any(
-            x
-            for x in associations_from_api
-            if x["SourceArn"] == artifact.artifact_arn and x["SourceType"] == "DataSet"
-        )
+        assert artifact.artifact_type == "DataSet"
 
 
 def test_training_job_arns(

tests/integ/sagemaker/lineage/test_model_artifact.py

@@ -12,11 +12,9 @@
 # language governing permissions and limitations under the License.
 """This module contains code to test SageMaker ``DatasetArtifact``"""
 from __future__ import absolute_import
-from tests.integ.sagemaker.lineage.helpers import traverse_graph_forward, traverse_graph_back
 
 
 def test_endpoints(
-    sagemaker_session,
     model_artifact_associated_endpoints,
     endpoint_deployment_action_obj,
     endpoint_context_obj,
@@ -32,44 +30,22 @@ def test_endpoints(
 
 def test_endpoint_contexts(
     static_model_artifact,
-    sagemaker_session,
 ):
     contexts_from_query = static_model_artifact.endpoint_contexts()
 
-    associations_from_api = traverse_graph_forward(
-        static_model_artifact.artifact_arn, sagemaker_session=sagemaker_session
-    )
-
     assert len(contexts_from_query) > 0
     for context in contexts_from_query:
-        # assert that the contexts from the query
-        # appear in the association list from the lineage API
-        assert any(
-            x
-            for x in associations_from_api
-            if x["DestinationArn"] == context.context_arn and x["DestinationType"] == "Endpoint"
-        )
+        assert context.context_type == "Endpoint"
 
 
 def test_dataset_artifacts(
     static_model_artifact,
-    sagemaker_session,
 ):
     artifacts_from_query = static_model_artifact.dataset_artifacts()
 
-    associations_from_api = traverse_graph_back(
-        static_model_artifact.artifact_arn, sagemaker_session=sagemaker_session
-    )
-
     assert len(artifacts_from_query) > 0
     for artifact in artifacts_from_query:
-        # assert that the artifacts from the query
-        # appear in the association list from the lineage API
-        assert any(
-            x
-            for x in associations_from_api
-            if x["SourceArn"] == artifact.artifact_arn and x["SourceType"] == "DataSet"
-        )
+        assert artifact.artifact_type == "DataSet"
 
 
 def test_training_job_arns(