test_get_visualization_elements added

Author: Yi-Ting Lee

src/sagemaker/lineage/query.py

@@ -364,7 +364,7 @@ def _get_visualization_elements(self):
         elements = {"nodes": verts, "edges": edges}
         return elements
 
-    def visualize(self):
+    def visualize(self, path="pyvisExample.html"):
         """Visualize lineage query result."""
         lineage_graph = {
             # nodes can have shape / color
@@ -398,7 +398,7 @@ def visualize(self):
 
         pyvis_vis = PyvisVisualizer(lineage_graph)
         elements = self._get_visualization_elements()
-        return pyvis_vis.render(elements=elements)
+        return pyvis_vis.render(elements=elements, path=path)
 
 
 class LineageFilter(object):

tests/integ/sagemaker/lineage/helpers.py

@@ -16,7 +16,8 @@
 import uuid
 from datetime import datetime
 import time
-
+import boto3
+from botocore.config import Config
 
 def name():
     return "lineage-integ-{}-{}".format(
@@ -78,3 +79,4 @@ def visit(arn, visited: set):
 
     ret = []
     return visit(start_arn, set())
+

tests/unit/sagemaker/lineage/test_query.py

@@ -524,3 +524,31 @@ def test_vertex_to_object_unconvertable(sagemaker_session):
 
     with pytest.raises(ValueError):
         vertex.to_lineage_object()
+
+
+def test_get_visualization_elements(sagemaker_session):
+    lineage_query = LineageQuery(sagemaker_session)
+    sagemaker_session.sagemaker_client.query_lineage.return_value = {
+        "Vertices": [
+            {"Arn": "arn1", "Type": "Endpoint", "LineageType": "Artifact"},
+            {"Arn": "arn2", "Type": "Model", "LineageType": "Context"},
+        ],
+        "Edges": [{"SourceArn": "arn1", "DestinationArn": "arn2", "AssociationType": "Produced"}],
+    }
+
+    query_response = lineage_query.query(
+        start_arns=["arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext"]
+    )
+
+    print(query_response)
+
+    elements = query_response._get_visualization_elements()
+
+    print(elements)
+
+    assert elements["nodes"][0] == ("arn1", "Endpoint", "Artifact", False)
+    assert elements["nodes"][1] == ("arn2", "Model", "Context", False)
+    assert elements["edges"][0] == ("arn1", "arn2", "Produced")
+
+    
+