def test_lineage_query_duplication(sagemaker_session):
lineage_query = LineageQuery(sagemaker_session)
sagemaker_session.sagemaker_client.query_lineage.return_value = {
"Vertices": [
{"Arn": "arn1", "Type": "Endpoint", "LineageType": "Artifact"},
{"Arn": "arn1", "Type": "Endpoint", "LineageType": "Artifact"},
{"Arn": "arn2", "Type": "Model", "LineageType": "Context"},
],
"Edges": [
{"SourceArn": "arn1", "DestinationArn": "arn2", "AssociationType": "Produced"},
{"SourceArn": "arn1", "DestinationArn": "arn2", "AssociationType": "Produced"},
],
}
response = lineage_query.query(
start_arns=["arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext"]
)
assert len(response.edges) == 1
assert response.edges[0].source_arn == "arn1"
assert response.edges[0].destination_arn == "arn2"
assert response.edges[0].association_type == "Produced"
assert len(response.vertices) == 2
assert response.vertices[0].arn == "arn1"
assert response.vertices[0].lineage_source == "Endpoint"
assert response.vertices[0].lineage_entity == "Artifact"
assert response.vertices[1].arn == "arn2"
assert response.vertices[1].lineage_source == "Model"
assert response.vertices[1].lineage_entity == "Context"
def test_lineage_query_cross_account_same_artifact(sagemaker_session):
lineage_query = LineageQuery(sagemaker_session)
sagemaker_session.sagemaker_client.query_lineage.return_value = {
"Vertices": [
{
"Arn": "arn:aws:sagemaker:us-east-2:012345678901:artifact/e1f29799189751939405b0f2b5b9d2a0",
"Type": "Endpoint",
"LineageType": "Artifact",
},
{
"Arn": "arn:aws:sagemaker:us-east-2:012345678902:artifact/e1f29799189751939405b0f2b5b9d2a0",
"Type": "Endpoint",
"LineageType": "Artifact",
},
],
"Edges": [
{
"SourceArn": "arn:aws:sagemaker:us-east-2:012345678901:artifact/e1f29799189751939405b0f2b5b9d2a0",
"DestinationArn": "arn:aws:sagemaker:us-east-2:012345678902:artifact/e1f29799189751939405b0f2b5b9d2a0",
"AssociationType": "SAME_AS",
},
{
"SourceArn": "arn:aws:sagemaker:us-east-2:012345678902:artifact/e1f29799189751939405b0f2b5b9d2a0",
"DestinationArn": "arn:aws:sagemaker:us-east-2:012345678901:artifact/e1f29799189751939405b0f2b5b9d2a0",
"AssociationType": "SAME_AS",
},
],
}
response = lineage_query.query(
start_arns=["arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext"]
)
assert len(response.edges) == 0
assert len(response.vertices) == 1
assert (
response.vertices[0].arn
== "arn:aws:sagemaker:us-east-2:012345678901:artifact/e1f29799189751939405b0f2b5b9d2a0"
)
assert response.vertices[0].lineage_source == "Endpoint"
assert response.vertices[0].lineage_entity == "Artifact"