diff --git a/tests/integ/sagemaker/lineage/conftest.py b/tests/integ/sagemaker/lineage/conftest.py index a03108e044..a3b2b844ad 100644 --- a/tests/integ/sagemaker/lineage/conftest.py +++ b/tests/integ/sagemaker/lineage/conftest.py @@ -663,7 +663,12 @@ def static_endpoint_context(sagemaker_session, static_pipeline_execution_arn): contexts = sagemaker_session.sagemaker_client.list_contexts(SourceUri=endpoint_arn)[ "ContextSummaries" ] - if len(contexts) != 1: + + logging.info(f"Found {len(contexts)} contexts associated with {endpoint_arn}") + for ctx in contexts: + logging.info(f'Found context "{ctx["ContextArn"]}"') + + if len(contexts) == 0: raise ( Exception( f"Got an unexpected number of Contexts for \ @@ -689,7 +694,12 @@ def static_model_package_group_context(sagemaker_session, static_pipeline_execut contexts = sagemaker_session.sagemaker_client.list_contexts(SourceUri=model_package_group_arn)[ "ContextSummaries" ] - if len(contexts) != 1: + + logging.info(f"Found {len(contexts)} contexts associated with {model_package_group_arn}") + for ctx in context: + logging.info(f'Found context "{ctx["ContextArn"]}"') + + if len(contexts) == 0: raise ( Exception( f"Got an unexpected number of Contexts for \ @@ -713,7 +723,12 @@ def static_model_artifact(sagemaker_session, static_pipeline_execution_arn): artifacts = sagemaker_session.sagemaker_client.list_artifacts(SourceUri=model_package_arn)[ "ArtifactSummaries" ] - if len(artifacts) != 1: + + logging.info(f"Found {len(artifacts)} artifacts associated with {model_package_arn}") + for art in artifacts: + logging.info(f'Found artifact {art["ArtifactArn"]}') + + if len(artifacts) == 0: raise ( Exception( f"Got an unexpected number of Artifacts for \