From b74e8610c6f3515b0c479b7802931c7f6af9defa Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Tue, 26 Jul 2022 16:36:53 -0700
Subject: [PATCH 01/11] test_lineage_visualize.py created

---
 tests/integ/sagemaker/lineage/helpers.py      | 56 +++++++++++++++++++
 .../lineage/test_lineage_visualize.py         | 51 +++++++++++++++++
 2 files changed, 107 insertions(+)
 create mode 100644 tests/integ/sagemaker/lineage/test_lineage_visualize.py

diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py
index fb71d1d88c..52eb44363c 100644
--- a/tests/integ/sagemaker/lineage/helpers.py
+++ b/tests/integ/sagemaker/lineage/helpers.py
@@ -16,6 +16,10 @@
 import uuid
 from datetime import datetime
 import time
+import boto3
+from sagemaker.lineage import association
+from sagemaker.lineage.artifact import Artifact
+from sagemaker.lineage.association import Association
 
 
 def name():
@@ -78,3 +82,55 @@ def visit(arn, visited: set):
 
     ret = []
     return visit(start_arn, set())
+
+
+class LineageResourceHelper:
+
+    def __init__(self):
+        self.client = boto3.client('sagemaker')
+        self.artifacts = []
+        self.associations = []
+
+    def create_artifact(self, artifact_name, artifact_type='Dataset'):
+        response = self.client.create_artifact(
+            ArtifactName=artifact_name,
+            Source={
+                'SourceUri': "Test-artifact-" + artifact_name,
+                'SourceTypes': [
+                    {
+                        'SourceIdType': 'S3ETag',
+                        'Value': 'Test-artifact-sourceId-value'
+                    },
+                ]
+            },
+            ArtifactType=artifact_type
+        )
+        self.artifacts.append(response['ArtifactArn'])
+
+        return response['ArtifactArn']
+
+    def create_association(self, source_arn, dest_arn, association_type='AssociatedWith'):
+        response = self.client.add_association(
+            SourceArn=source_arn,
+            DestinationArn=dest_arn,
+            AssociationType=association_type
+        )
+        if "SourceArn" in response.keys():
+            self.associations.append((source_arn, dest_arn))
+            return True
+        else:
+            return False
+
+    def clean_all(self):
+        for source, dest in self.associations:
+            try:
+                self.client.delete_association(
+                    SourceArn=source,
+                    DestinationArn=dest
+                )
+                time.sleep(2)
+            except(e):
+                print("skipped " + str(e))
+
+        for artifact_arn in self.artifacts:
+            self.client.delete_artifact(ArtifactArn=artifact_arn)
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
new file mode 100644
index 0000000000..cb68021da5
--- /dev/null
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -0,0 +1,51 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module contains code to test SageMaker ``LineageQueryResult.visualize()``"""
+
+import datetime
+import logging
+import time
+
+import pytest
+
+import sagemaker.lineage.query
+
+from tests.integ.sagemaker.lineage.helpers import name, names, retry, LineageResourceHelper
+
+def test_LineageResourceHelper():
+    lineage_resource_helper = LineageResourceHelper()
+    art1 = lineage_resource_helper.create_artifact(artifact_name=name())
+    art2 = lineage_resource_helper.create_artifact(artifact_name=name())
+    lineage_resource_helper.create_association(source_arn=art1, dest_arn=art2)
+    lineage_resource_helper.clean_all()
+
+def test_wide_graphs(sagemaker_session):
+    lineage_resource_helper = LineageResourceHelper()
+    art_root = lineage_resource_helper.create_artifact(artifact_name=name())
+    try:
+        for i in range(10):
+            art = lineage_resource_helper.create_artifact(artifact_name=name())
+            lineage_resource_helper.create_association(source_arn=art_root, dest_arn=art)
+            time.sleep(0.1)
+    except(e):
+        lineage_resource_helper.clean_all()
+
+    try:
+        lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
+        result = lq.query(start_arns=[art_root])
+        print(result)
+    except(e):
+        lineage_resource_helper.clean_all()
+
+    lineage_resource_helper.clean_all()
+

From 578acdf8dd620180370ae8ab5e68579ca663cfab Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Wed, 27 Jul 2022 14:47:06 -0700
Subject: [PATCH 02/11] pyvis import issue on running lineage test

---
 requirements/extras/test_requirements.txt     |  1 +
 tests/integ/sagemaker/lineage/helpers.py      | 35 +++++++------------
 .../lineage/test_lineage_visualize.py         | 15 ++++----
 3 files changed, 23 insertions(+), 28 deletions(-)

diff --git a/requirements/extras/test_requirements.txt b/requirements/extras/test_requirements.txt
index 2247394441..7197147032 100644
--- a/requirements/extras/test_requirements.txt
+++ b/requirements/extras/test_requirements.txt
@@ -18,3 +18,4 @@ fabric==2.6.0
 requests==2.27.1
 sagemaker-experiments==0.1.35
 Jinja2==3.0.3
+pyvis==0.2.1
diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py
index 52eb44363c..7b317c5705 100644
--- a/tests/integ/sagemaker/lineage/helpers.py
+++ b/tests/integ/sagemaker/lineage/helpers.py
@@ -85,35 +85,29 @@ def visit(arn, visited: set):
 
 
 class LineageResourceHelper:
-
     def __init__(self):
-        self.client = boto3.client('sagemaker')
+        self.client = boto3.client("sagemaker")
         self.artifacts = []
         self.associations = []
 
-    def create_artifact(self, artifact_name, artifact_type='Dataset'):
+    def create_artifact(self, artifact_name, artifact_type="Dataset"):
         response = self.client.create_artifact(
             ArtifactName=artifact_name,
             Source={
-                'SourceUri': "Test-artifact-" + artifact_name,
-                'SourceTypes': [
-                    {
-                        'SourceIdType': 'S3ETag',
-                        'Value': 'Test-artifact-sourceId-value'
-                    },
-                ]
+                "SourceUri": "Test-artifact-" + artifact_name,
+                "SourceTypes": [
+                    {"SourceIdType": "S3ETag", "Value": "Test-artifact-sourceId-value"},
+                ],
             },
-            ArtifactType=artifact_type
+            ArtifactType=artifact_type,
         )
-        self.artifacts.append(response['ArtifactArn'])
+        self.artifacts.append(response["ArtifactArn"])
 
-        return response['ArtifactArn']
+        return response["ArtifactArn"]
 
-    def create_association(self, source_arn, dest_arn, association_type='AssociatedWith'):
+    def create_association(self, source_arn, dest_arn, association_type="AssociatedWith"):
         response = self.client.add_association(
-            SourceArn=source_arn,
-            DestinationArn=dest_arn,
-            AssociationType=association_type
+            SourceArn=source_arn, DestinationArn=dest_arn, AssociationType=association_type
         )
         if "SourceArn" in response.keys():
             self.associations.append((source_arn, dest_arn))
@@ -124,12 +118,9 @@ def create_association(self, source_arn, dest_arn, association_type='AssociatedW
     def clean_all(self):
         for source, dest in self.associations:
             try:
-                self.client.delete_association(
-                    SourceArn=source,
-                    DestinationArn=dest
-                )
+                self.client.delete_association(SourceArn=source, DestinationArn=dest)
                 time.sleep(2)
-            except(e):
+            except (e):
                 print("skipped " + str(e))
 
         for artifact_arn in self.artifacts:
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index cb68021da5..eede2eaf9f 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -22,6 +22,7 @@
 
 from tests.integ.sagemaker.lineage.helpers import name, names, retry, LineageResourceHelper
 
+
 def test_LineageResourceHelper():
     lineage_resource_helper = LineageResourceHelper()
     art1 = lineage_resource_helper.create_artifact(artifact_name=name())
@@ -29,6 +30,7 @@ def test_LineageResourceHelper():
     lineage_resource_helper.create_association(source_arn=art1, dest_arn=art2)
     lineage_resource_helper.clean_all()
 
+
 def test_wide_graphs(sagemaker_session):
     lineage_resource_helper = LineageResourceHelper()
     art_root = lineage_resource_helper.create_artifact(artifact_name=name())
@@ -36,16 +38,17 @@ def test_wide_graphs(sagemaker_session):
         for i in range(10):
             art = lineage_resource_helper.create_artifact(artifact_name=name())
             lineage_resource_helper.create_association(source_arn=art_root, dest_arn=art)
-            time.sleep(0.1)
-    except(e):
+            time.sleep(0.2)
+    except Exception as e:
+        print(e)
         lineage_resource_helper.clean_all()
 
     try:
         lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
-        result = lq.query(start_arns=[art_root])
-        print(result)
-    except(e):
+        lq_result = lq.query(start_arns=[art_root])
+        lq_result.visualize()
+    except Exception as e:
+        print(e)
         lineage_resource_helper.clean_all()
 
     lineage_resource_helper.clean_all()
-

From 9171e4e6afecd2d7b19e251f2c361704c0c4fcb4 Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Wed, 27 Jul 2022 15:02:14 -0700
Subject: [PATCH 03/11] import visualization modules using get_module()

---
 src/sagemaker/lineage/query.py | 28 +++++-----------------------
 1 file changed, 5 insertions(+), 23 deletions(-)

diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py
index b96880209b..95c5f6fa22 100644
--- a/src/sagemaker/lineage/query.py
+++ b/src/sagemaker/lineage/query.py
@@ -17,7 +17,7 @@
 from enum import Enum
 from typing import Optional, Union, List, Dict
 
-from sagemaker.lineage._utils import get_resource_name_from_arn
+from sagemaker.lineage._utils import get_resource_name_from_arn, get_module
 
 
 class LineageEntityEnum(Enum):
@@ -208,7 +208,6 @@ def __init__(self, graph_styles):
         """Init for PyvisVisualizer."""
         # import visualization packages
         (
-            self.pyvis,
             self.Network,
             self.Options,
         ) = self._import_visual_modules()
@@ -217,28 +216,11 @@ def __init__(self, graph_styles):
 
     def _import_visual_modules(self):
         """Import modules needed for visualization."""
-        try:
-            import pyvis
-        except ImportError as e:
-            print(e)
-            print("Try: pip install pyvis")
-            raise
-
-        try:
-            from pyvis.network import Network
-        except ImportError as e:
-            print(e)
-            print("Try: pip install pyvis")
-            raise
-
-        try:
-            from pyvis.options import Options
-        except ImportError as e:
-            print(e)
-            print("Try: pip install pyvis")
-            raise
+        get_module("pyvis")
+        from pyvis.network import Network
+        from pyvis.options import Options
 
-        return pyvis, Network, Options
+        return Network, Options
 
     def _get_options(self):
         """Get pyvis graph options."""

From d679f4a015a16bfb1d950c93158e7681ef5d0bbb Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Thu, 28 Jul 2022 09:51:40 -0700
Subject: [PATCH 04/11] test_wide_graphs added

---
 .gitignore                                       |  3 ++-
 requirements/extras/local_requirements.txt       |  2 +-
 tests/integ/sagemaker/lineage/helpers.py         | 13 +++++++++----
 .../sagemaker/lineage/test_lineage_visualize.py  | 16 +++++++++++-----
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/.gitignore b/.gitignore
index 5b496055e9..7a61658c20 100644
--- a/.gitignore
+++ b/.gitignore
@@ -29,4 +29,5 @@ venv/
 env/
 .vscode/
 **/tmp
-.python-version
\ No newline at end of file
+.python-version
+*.html
\ No newline at end of file
diff --git a/requirements/extras/local_requirements.txt b/requirements/extras/local_requirements.txt
index 17512c3388..439f42010c 100644
--- a/requirements/extras/local_requirements.txt
+++ b/requirements/extras/local_requirements.txt
@@ -1,4 +1,4 @@
 urllib3==1.26.8
 docker-compose==1.29.2
 docker~=5.0.0
-PyYAML==5.4.1
+PyYAML==5.4.1
\ No newline at end of file
diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py
index 7b317c5705..361019d9b4 100644
--- a/tests/integ/sagemaker/lineage/helpers.py
+++ b/tests/integ/sagemaker/lineage/helpers.py
@@ -17,6 +17,7 @@
 from datetime import datetime
 import time
 import boto3
+from botocore.config import Config
 from sagemaker.lineage import association
 from sagemaker.lineage.artifact import Artifact
 from sagemaker.lineage.association import Association
@@ -86,7 +87,7 @@ def visit(arn, visited: set):
 
 class LineageResourceHelper:
     def __init__(self):
-        self.client = boto3.client("sagemaker")
+        self.client = boto3.client("sagemaker", config=Config(connect_timeout=5, read_timeout=60, retries={'max_attempts': 20}))
         self.artifacts = []
         self.associations = []
 
@@ -119,9 +120,13 @@ def clean_all(self):
         for source, dest in self.associations:
             try:
                 self.client.delete_association(SourceArn=source, DestinationArn=dest)
-                time.sleep(2)
-            except (e):
+                time.sleep(0.5)
+            except Exception as e:
                 print("skipped " + str(e))
 
         for artifact_arn in self.artifacts:
-            self.client.delete_artifact(ArtifactArn=artifact_arn)
+            try:
+                self.client.delete_artifact(ArtifactArn=artifact_arn)
+                time.sleep(0.5)
+            except Exception as e:
+                print("skipped " + str(e))
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index eede2eaf9f..3c84ae057b 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -25,23 +25,28 @@
 
 def test_LineageResourceHelper():
     lineage_resource_helper = LineageResourceHelper()
-    art1 = lineage_resource_helper.create_artifact(artifact_name=name())
-    art2 = lineage_resource_helper.create_artifact(artifact_name=name())
-    lineage_resource_helper.create_association(source_arn=art1, dest_arn=art2)
-    lineage_resource_helper.clean_all()
+    try:
+        art1 = lineage_resource_helper.create_artifact(artifact_name=name())
+        art2 = lineage_resource_helper.create_artifact(artifact_name=name())
+        lineage_resource_helper.create_association(source_arn=art1, dest_arn=art2)
+        lineage_resource_helper.clean_all()
+    except Exception as e:
+        print(e)
+        assert False
 
 
 def test_wide_graphs(sagemaker_session):
     lineage_resource_helper = LineageResourceHelper()
     art_root = lineage_resource_helper.create_artifact(artifact_name=name())
     try:
-        for i in range(10):
+        for i in range(500):
             art = lineage_resource_helper.create_artifact(artifact_name=name())
             lineage_resource_helper.create_association(source_arn=art_root, dest_arn=art)
             time.sleep(0.2)
     except Exception as e:
         print(e)
         lineage_resource_helper.clean_all()
+        assert False
 
     try:
         lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
@@ -50,5 +55,6 @@ def test_wide_graphs(sagemaker_session):
     except Exception as e:
         print(e)
         lineage_resource_helper.clean_all()
+        assert False
 
     lineage_resource_helper.clean_all()

From 5c89c0931d9ce9c0b1bbfaca443fb5fb98795602 Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Thu, 28 Jul 2022 13:46:31 -0700
Subject: [PATCH 05/11] long graph visualize test added

---
 src/sagemaker/lineage/query.py                |  5 +-
 .../lineage/test_lineage_visualize.py         | 52 ++++++++++++++++---
 2 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py
index 95c5f6fa22..a5c8c0bfae 100644
--- a/src/sagemaker/lineage/query.py
+++ b/src/sagemaker/lineage/query.py
@@ -274,7 +274,6 @@ def render(self, elements, path="pyvisExample.html"):
 
         # add edges to graph
         for e in elements["edges"]:
-            print(e)
             net.add_edge(e[0], e[1], title=e[2])
 
         return net.show(path)
@@ -366,7 +365,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
@@ -400,7 +399,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):
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index 3c84ae057b..0fb1c04f43 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -19,11 +19,12 @@
 import pytest
 
 import sagemaker.lineage.query
-
+from sagemaker.lineage.query import LineageQueryDirectionEnum
 from tests.integ.sagemaker.lineage.helpers import name, names, retry, LineageResourceHelper
 
 
 def test_LineageResourceHelper():
+    # check if LineageResourceHelper works properly
     lineage_resource_helper = LineageResourceHelper()
     try:
         art1 = lineage_resource_helper.create_artifact(artifact_name=name())
@@ -35,13 +36,19 @@ def test_LineageResourceHelper():
         assert False
 
 
-def test_wide_graphs(sagemaker_session):
+def test_wide_graph_visualize(sagemaker_session):
     lineage_resource_helper = LineageResourceHelper()
-    art_root = lineage_resource_helper.create_artifact(artifact_name=name())
+    wide_graph_root_arn = lineage_resource_helper.create_artifact(artifact_name=name())
+
+    # create wide graph
+    # Artifact ----> Artifact
+    #        \ \ \-> Artifact
+    #         \ \--> Artifact
+    #          \--->  ...
     try:
-        for i in range(500):
-            art = lineage_resource_helper.create_artifact(artifact_name=name())
-            lineage_resource_helper.create_association(source_arn=art_root, dest_arn=art)
+        for i in range(3):
+            artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
+            lineage_resource_helper.create_association(source_arn=wide_graph_root_arn, dest_arn=artifact_arn)
             time.sleep(0.2)
     except Exception as e:
         print(e)
@@ -50,11 +57,40 @@ def test_wide_graphs(sagemaker_session):
 
     try:
         lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
-        lq_result = lq.query(start_arns=[art_root])
-        lq_result.visualize()
+        lq_result = lq.query(start_arns=[wide_graph_root_arn])
+        lq_result.visualize(path="wideGraph.html")
     except Exception as e:
         print(e)
         lineage_resource_helper.clean_all()
         assert False
 
     lineage_resource_helper.clean_all()
+
+def test_long_graph_visualize(sagemaker_session):
+    lineage_resource_helper = LineageResourceHelper()
+    long_graph_root_arn = lineage_resource_helper.create_artifact(artifact_name=name())
+    last_arn = long_graph_root_arn
+
+    # create long graph
+    # Artifact -> Artifact -> ... -> Artifact
+    try:
+        for i in range(20):
+            new_artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
+            lineage_resource_helper.create_association(source_arn=last_arn, dest_arn=new_artifact_arn)
+            last_arn = new_artifact_arn
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    try:
+        lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
+        lq_result = lq.query(start_arns=[long_graph_root_arn], direction=LineageQueryDirectionEnum.DESCENDANTS)
+        # max depth = 10 -> graph rendered only has length of ten (in DESCENDANTS direction)
+        lq_result.visualize(path="longGraph.html")
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    lineage_resource_helper.clean_all()
\ No newline at end of file

From fa7b9dd752b00740060a22ae95be5f87597f35df Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Fri, 29 Jul 2022 14:06:09 -0700
Subject: [PATCH 06/11] create context & action added to helper

---
 tests/integ/sagemaker/lineage/helpers.py      | 51 ++++++++++++++++++-
 .../lineage/test_lineage_visualize.py         | 14 ++---
 2 files changed, 56 insertions(+), 9 deletions(-)

diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py
index 361019d9b4..4e73537584 100644
--- a/tests/integ/sagemaker/lineage/helpers.py
+++ b/tests/integ/sagemaker/lineage/helpers.py
@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 """This module contains helper methods for tests of SageMaker Lineage"""
 from __future__ import absolute_import
+from urllib import response
 
 import uuid
 from datetime import datetime
@@ -86,9 +87,12 @@ def visit(arn, visited: set):
 
 
 class LineageResourceHelper:
-    def __init__(self):
-        self.client = boto3.client("sagemaker", config=Config(connect_timeout=5, read_timeout=60, retries={'max_attempts': 20}))
+    def __init__(self, sagemaker_session):
+        self.client = sagemaker_session.sagemaker_client
         self.artifacts = []
+        self.actions = []
+        self.contexts = []
+        self.trialComponents = []
         self.associations = []
 
     def create_artifact(self, artifact_name, artifact_type="Dataset"):
@@ -106,6 +110,42 @@ def create_artifact(self, artifact_name, artifact_type="Dataset"):
 
         return response["ArtifactArn"]
 
+    def create_action(self, action_name, action_type="ModelDeployment"):
+        response = self.client.create_action(
+            ActionName=action_name,
+            Source={
+                "SourceUri": "Test-action-" + action_name,
+                "SourceTypes": [
+                    {"SourceIdType": "S3ETag", "Value": "Test-action-sourceId-value"},
+                ],
+            },
+            ActionType=action_type
+        )
+        self.actions.append(response["ActionArn"])
+
+        return response["ActionArn"]
+
+    def create_context(self, context_name, context_type="Endpoint"):
+        response = self.client.create_context(
+            ContextName=context_name,
+            Source={
+                "SourceUri": "Test-context-" + context_name,
+                "SourceTypes": [
+                    {"SourceIdType": "S3ETag", "Value": "Test-context-sourceId-value"},
+                ],
+            },
+            ContextType=context_type
+        )
+        self.contexts.append(response["ContextArn"])
+
+        return response["ContextArn"]
+
+    def create_trialComponent(self, trialComponent_name, trialComponent_type="TrainingJob"):
+        response = self.client.create_trial_component(
+            TrialComponentName=trialComponent_name,
+            
+        )
+
     def create_association(self, source_arn, dest_arn, association_type="AssociatedWith"):
         response = self.client.add_association(
             SourceArn=source_arn, DestinationArn=dest_arn, AssociationType=association_type
@@ -130,3 +170,10 @@ def clean_all(self):
                 time.sleep(0.5)
             except Exception as e:
                 print("skipped " + str(e))
+
+        for action_arn in self.actions:
+            try:
+                self.client.delete_action(ActionArn=action_arn)
+                time.sleep(0.5)
+            except Exception as e:
+                print("skipped " + str(e))
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index 0fb1c04f43..5f12386392 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -23,9 +23,9 @@
 from tests.integ.sagemaker.lineage.helpers import name, names, retry, LineageResourceHelper
 
 
-def test_LineageResourceHelper():
+def test_LineageResourceHelper(sagemaker_session):
     # check if LineageResourceHelper works properly
-    lineage_resource_helper = LineageResourceHelper()
+    lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
     try:
         art1 = lineage_resource_helper.create_artifact(artifact_name=name())
         art2 = lineage_resource_helper.create_artifact(artifact_name=name())
@@ -35,9 +35,9 @@ def test_LineageResourceHelper():
         print(e)
         assert False
 
-
+@pytest.mark.skip("visualizer load test")
 def test_wide_graph_visualize(sagemaker_session):
-    lineage_resource_helper = LineageResourceHelper()
+    lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
     wide_graph_root_arn = lineage_resource_helper.create_artifact(artifact_name=name())
 
     # create wide graph
@@ -46,10 +46,9 @@ def test_wide_graph_visualize(sagemaker_session):
     #         \ \--> Artifact
     #          \--->  ...
     try:
-        for i in range(3):
+        for i in range(500):
             artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
             lineage_resource_helper.create_association(source_arn=wide_graph_root_arn, dest_arn=artifact_arn)
-            time.sleep(0.2)
     except Exception as e:
         print(e)
         lineage_resource_helper.clean_all()
@@ -66,8 +65,9 @@ def test_wide_graph_visualize(sagemaker_session):
 
     lineage_resource_helper.clean_all()
 
+@pytest.mark.skip("visualizer load test")
 def test_long_graph_visualize(sagemaker_session):
-    lineage_resource_helper = LineageResourceHelper()
+    lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
     long_graph_root_arn = lineage_resource_helper.create_artifact(artifact_name=name())
     last_arn = long_graph_root_arn
 

From 83964da9b479d9ff4f034ce179f95b3fa56554ce Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Fri, 29 Jul 2022 14:22:21 -0700
Subject: [PATCH 07/11] resolve conflict with master branch

---
 src/sagemaker/lineage/query.py | 42 ++++++++++++++++------------------
 1 file changed, 20 insertions(+), 22 deletions(-)

diff --git a/src/sagemaker/lineage/query.py b/src/sagemaker/lineage/query.py
index a5c8c0bfae..f5ff1edef1 100644
--- a/src/sagemaker/lineage/query.py
+++ b/src/sagemaker/lineage/query.py
@@ -214,17 +214,7 @@ def __init__(self, graph_styles):
 
         self.graph_styles = graph_styles
 
-    def _import_visual_modules(self):
-        """Import modules needed for visualization."""
-        get_module("pyvis")
-        from pyvis.network import Network
-        from pyvis.options import Options
-
-        return Network, Options
-
-    def _get_options(self):
-        """Get pyvis graph options."""
-        options = """
+        self._options = """
             var options = {
             "configure":{
                 "enabled": false
@@ -253,28 +243,36 @@ def _get_options(self):
             }
         }
             """
-        return options
 
-    def _node_color(self, n):
+    def _import_visual_modules(self):
+        """Import modules needed for visualization."""
+        get_module("pyvis")
+        from pyvis.network import Network
+        from pyvis.options import Options
+
+        return Network, Options
+
+    def _node_color(self, entity):
         """Return node color by background-color specified in graph styles."""
-        return self.graph_styles[n[2]]["style"]["background-color"]
+        return self.graph_styles[entity]["style"]["background-color"]
 
     def render(self, elements, path="pyvisExample.html"):
         """Render graph for lineage query result."""
         net = self.Network(height="500px", width="100%", notebook=True, directed=True)
-        options = self._get_options()
-        net.set_options(options)
+        net.set_options(self._options)
 
         # add nodes to graph
-        for n in elements["nodes"]:
-            if n[3]:  # startarn
-                net.add_node(n[0], label=n[1], title=n[2], color=self._node_color(n), shape="star")
+        for arn, source, entity, is_start_arn in elements["nodes"]:
+            if is_start_arn:  # startarn
+                net.add_node(
+                    arn, label=source, title=entity, color=self._node_color(entity), shape="star"
+                )
             else:
-                net.add_node(n[0], label=n[1], title=n[2], color=self._node_color(n))
+                net.add_node(arn, label=source, title=entity, color=self._node_color(entity))
 
         # add edges to graph
-        for e in elements["edges"]:
-            net.add_edge(e[0], e[1], title=e[2])
+        for src, dest, asso_type in elements["edges"]:
+            net.add_edge(src, dest, title=asso_type)
 
         return net.show(path)
 

From c66edf87243c7b2819ea0681aef16576cdded30f Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Mon, 1 Aug 2022 14:17:32 -0700
Subject: [PATCH 08/11] change: add queryLineageResult visualizer load test &
 integ test

---
 tests/integ/sagemaker/lineage/helpers.py      |  34 ++-
 .../lineage/test_lineage_visualize.py         | 194 +++++++++++++++++-
 2 files changed, 197 insertions(+), 31 deletions(-)

diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py
index 4e73537584..0c40bbac91 100644
--- a/tests/integ/sagemaker/lineage/helpers.py
+++ b/tests/integ/sagemaker/lineage/helpers.py
@@ -12,16 +12,10 @@
 # language governing permissions and limitations under the License.
 """This module contains helper methods for tests of SageMaker Lineage"""
 from __future__ import absolute_import
-from urllib import response
 
 import uuid
 from datetime import datetime
 import time
-import boto3
-from botocore.config import Config
-from sagemaker.lineage import association
-from sagemaker.lineage.artifact import Artifact
-from sagemaker.lineage.association import Association
 
 
 def name():
@@ -92,7 +86,6 @@ def __init__(self, sagemaker_session):
         self.artifacts = []
         self.actions = []
         self.contexts = []
-        self.trialComponents = []
         self.associations = []
 
     def create_artifact(self, artifact_name, artifact_type="Dataset"):
@@ -115,11 +108,10 @@ def create_action(self, action_name, action_type="ModelDeployment"):
             ActionName=action_name,
             Source={
                 "SourceUri": "Test-action-" + action_name,
-                "SourceTypes": [
-                    {"SourceIdType": "S3ETag", "Value": "Test-action-sourceId-value"},
-                ],
+                "SourceType": "S3ETag",
+                "SourceId": "Test-action-sourceId-value",
             },
-            ActionType=action_type
+            ActionType=action_type,
         )
         self.actions.append(response["ActionArn"])
 
@@ -130,22 +122,15 @@ def create_context(self, context_name, context_type="Endpoint"):
             ContextName=context_name,
             Source={
                 "SourceUri": "Test-context-" + context_name,
-                "SourceTypes": [
-                    {"SourceIdType": "S3ETag", "Value": "Test-context-sourceId-value"},
-                ],
+                "SourceType": "S3ETag",
+                "SourceId": "Test-context-sourceId-value",
             },
-            ContextType=context_type
+            ContextType=context_type,
         )
         self.contexts.append(response["ContextArn"])
 
         return response["ContextArn"]
 
-    def create_trialComponent(self, trialComponent_name, trialComponent_type="TrainingJob"):
-        response = self.client.create_trial_component(
-            TrialComponentName=trialComponent_name,
-            
-        )
-
     def create_association(self, source_arn, dest_arn, association_type="AssociatedWith"):
         response = self.client.add_association(
             SourceArn=source_arn, DestinationArn=dest_arn, AssociationType=association_type
@@ -177,3 +162,10 @@ def clean_all(self):
                 time.sleep(0.5)
             except Exception as e:
                 print("skipped " + str(e))
+
+        for context_arn in self.contexts:
+            try:
+                self.client.delete_context(ContextArn=context_arn)
+                time.sleep(0.5)
+            except Exception as e:
+                print("skipped " + str(e))
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index 5f12386392..274b18d503 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -11,16 +11,15 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """This module contains code to test SageMaker ``LineageQueryResult.visualize()``"""
-
-import datetime
-import logging
+from __future__ import absolute_import
 import time
+import json
 
 import pytest
 
 import sagemaker.lineage.query
 from sagemaker.lineage.query import LineageQueryDirectionEnum
-from tests.integ.sagemaker.lineage.helpers import name, names, retry, LineageResourceHelper
+from tests.integ.sagemaker.lineage.helpers import name, LineageResourceHelper
 
 
 def test_LineageResourceHelper(sagemaker_session):
@@ -35,6 +34,7 @@ def test_LineageResourceHelper(sagemaker_session):
         print(e)
         assert False
 
+
 @pytest.mark.skip("visualizer load test")
 def test_wide_graph_visualize(sagemaker_session):
     lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
@@ -46,9 +46,11 @@ def test_wide_graph_visualize(sagemaker_session):
     #         \ \--> Artifact
     #          \--->  ...
     try:
-        for i in range(500):
+        for i in range(10):
             artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
-            lineage_resource_helper.create_association(source_arn=wide_graph_root_arn, dest_arn=artifact_arn)
+            lineage_resource_helper.create_association(
+                source_arn=wide_graph_root_arn, dest_arn=artifact_arn
+            )
     except Exception as e:
         print(e)
         lineage_resource_helper.clean_all()
@@ -65,6 +67,7 @@ def test_wide_graph_visualize(sagemaker_session):
 
     lineage_resource_helper.clean_all()
 
+
 @pytest.mark.skip("visualizer load test")
 def test_long_graph_visualize(sagemaker_session):
     lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
@@ -74,9 +77,11 @@ def test_long_graph_visualize(sagemaker_session):
     # create long graph
     # Artifact -> Artifact -> ... -> Artifact
     try:
-        for i in range(20):
+        for i in range(10):
             new_artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
-            lineage_resource_helper.create_association(source_arn=last_arn, dest_arn=new_artifact_arn)
+            lineage_resource_helper.create_association(
+                source_arn=last_arn, dest_arn=new_artifact_arn
+            )
             last_arn = new_artifact_arn
     except Exception as e:
         print(e)
@@ -85,7 +90,9 @@ def test_long_graph_visualize(sagemaker_session):
 
     try:
         lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
-        lq_result = lq.query(start_arns=[long_graph_root_arn], direction=LineageQueryDirectionEnum.DESCENDANTS)
+        lq_result = lq.query(
+            start_arns=[long_graph_root_arn], direction=LineageQueryDirectionEnum.DESCENDANTS
+        )
         # max depth = 10 -> graph rendered only has length of ten (in DESCENDANTS direction)
         lq_result.visualize(path="longGraph.html")
     except Exception as e:
@@ -93,4 +100,171 @@ def test_long_graph_visualize(sagemaker_session):
         lineage_resource_helper.clean_all()
         assert False
 
-    lineage_resource_helper.clean_all()
\ No newline at end of file
+    lineage_resource_helper.clean_all()
+
+
+def test_graph_visualize(sagemaker_session):
+    lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
+
+    # create lineage data
+    # image artifact ------> model artifact(startarn) -> model deploy action -> endpoint context
+    #                    /->
+    # dataset artifact -/
+    try:
+        graph_startarn = lineage_resource_helper.create_artifact(
+            artifact_name=name(), artifact_type="Model"
+        )
+        image_artifact = lineage_resource_helper.create_artifact(
+            artifact_name=name(), artifact_type="Image"
+        )
+        lineage_resource_helper.create_association(
+            source_arn=image_artifact, dest_arn=graph_startarn, association_type="ContributedTo"
+        )
+        dataset_artifact = lineage_resource_helper.create_artifact(
+            artifact_name=name(), artifact_type="DataSet"
+        )
+        lineage_resource_helper.create_association(
+            source_arn=dataset_artifact, dest_arn=graph_startarn, association_type="AssociatedWith"
+        )
+        modeldeploy_action = lineage_resource_helper.create_action(
+            action_name=name(), action_type="ModelDeploy"
+        )
+        lineage_resource_helper.create_association(
+            source_arn=graph_startarn, dest_arn=modeldeploy_action, association_type="ContributedTo"
+        )
+        endpoint_context = lineage_resource_helper.create_context(
+            context_name=name(), context_type="Endpoint"
+        )
+        lineage_resource_helper.create_association(
+            source_arn=modeldeploy_action,
+            dest_arn=endpoint_context,
+            association_type="AssociatedWith",
+        )
+        time.sleep(1)
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    # visualize
+    try:
+        lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
+        lq_result = lq.query(start_arns=[graph_startarn])
+        lq_result.visualize(path="testGraph.html")
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    # check generated graph info
+    try:
+        fo = open("testGraph.html", "r")
+        lines = fo.readlines()
+        for line in lines:
+            if "nodes = " in line:
+                node = line
+            if "edges = " in line:
+                edge = line
+
+        # extract node data
+        start = node.find("[")
+        end = node.find("]")
+        res = node[start + 1 : end].split("}, ")
+        res = [i + "}" for i in res]
+        res[-1] = res[-1][:-1]
+        node_dict = [json.loads(i) for i in res]
+
+        # extract edge data
+        start = edge.find("[")
+        end = edge.find("]")
+        res = edge[start + 1 : end].split("}, ")
+        res = [i + "}" for i in res]
+        res[-1] = res[-1][:-1]
+        edge_dict = [json.loads(i) for i in res]
+
+        # check node number
+        assert len(node_dict) == 5
+
+        # check startarn
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == graph_startarn
+        )
+        assert found_value["color"] == "#146eb4"
+        assert found_value["label"] == "Model"
+        assert found_value["shape"] == "star"
+        assert found_value["title"] == "Artifact"
+
+        # check image artifact
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == image_artifact
+        )
+        assert found_value["color"] == "#146eb4"
+        assert found_value["label"] == "Image"
+        assert found_value["shape"] == "dot"
+        assert found_value["title"] == "Artifact"
+
+        # check dataset artifact
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == dataset_artifact
+        )
+        assert found_value["color"] == "#146eb4"
+        assert found_value["label"] == "DataSet"
+        assert found_value["shape"] == "dot"
+        assert found_value["title"] == "Artifact"
+
+        # check modeldeploy action
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == modeldeploy_action
+        )
+        assert found_value["color"] == "#88c396"
+        assert found_value["label"] == "ModelDeploy"
+        assert found_value["shape"] == "dot"
+        assert found_value["title"] == "Action"
+
+        # check endpoint context
+        found_value = next(
+            dictionary for dictionary in node_dict if dictionary["id"] == endpoint_context
+        )
+        assert found_value["color"] == "#ff9900"
+        assert found_value["label"] == "Endpoint"
+        assert found_value["shape"] == "dot"
+        assert found_value["title"] == "Context"
+
+        # check edge number
+        assert len(edge_dict) == 4
+
+        # check image_artifact ->  model_artifact(startarn) edge
+        found_value = next(
+            dictionary for dictionary in edge_dict if dictionary["from"] == image_artifact
+        )
+        assert found_value["to"] == graph_startarn
+        assert found_value["title"] == "ContributedTo"
+
+        # check dataset_artifact ->  model_artifact(startarn) edge
+        found_value = next(
+            dictionary for dictionary in edge_dict if dictionary["from"] == dataset_artifact
+        )
+        assert found_value["to"] == graph_startarn
+        assert found_value["title"] == "AssociatedWith"
+
+        # check model_artifact(startarn) ->  modeldeploy_action edge
+        found_value = next(
+            dictionary for dictionary in edge_dict if dictionary["from"] == graph_startarn
+        )
+        assert found_value["to"] == modeldeploy_action
+        assert found_value["title"] == "ContributedTo"
+
+        # check modeldeploy_action ->  endpoint_context edge
+        found_value = next(
+            dictionary for dictionary in edge_dict if dictionary["from"] == modeldeploy_action
+        )
+        assert found_value["to"] == endpoint_context
+        assert found_value["title"] == "AssociatedWith"
+
+    except Exception as e:
+        print(e)
+        lineage_resource_helper.clean_all()
+        assert False
+
+    # clean lineage data
+    lineage_resource_helper.clean_all()

From 7cec38d2c88aafb402c23818ce3bb3b6e55ce505 Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Mon, 1 Aug 2022 15:59:04 -0700
Subject: [PATCH 09/11] remove generated graph html file after visualize integ
 test

---
 tests/integ/sagemaker/lineage/test_lineage_visualize.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index 274b18d503..1ba87d6a3c 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -14,6 +14,7 @@
 from __future__ import absolute_import
 import time
 import json
+import os
 
 import pytest
 
@@ -264,7 +265,10 @@ def test_graph_visualize(sagemaker_session):
     except Exception as e:
         print(e)
         lineage_resource_helper.clean_all()
+        os.remove("testGraph.html")
         assert False
 
+    # delete generated test graph
+    os.remove("testGraph.html")
     # clean lineage data
     lineage_resource_helper.clean_all()

From c1169be2c0ea443a85bcc2325fc4b46df1393ce7 Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Wed, 3 Aug 2022 12:03:42 -0700
Subject: [PATCH 10/11] validation logic clean on integ tests

---
 tests/integ/sagemaker/lineage/conftest.py     |  15 ++
 tests/integ/sagemaker/lineage/helpers.py      |   4 -
 .../lineage/test_lineage_visualize.py         | 202 +++++++-----------
 3 files changed, 95 insertions(+), 126 deletions(-)

diff --git a/tests/integ/sagemaker/lineage/conftest.py b/tests/integ/sagemaker/lineage/conftest.py
index 5e201eef42..7450cc5935 100644
--- a/tests/integ/sagemaker/lineage/conftest.py
+++ b/tests/integ/sagemaker/lineage/conftest.py
@@ -19,6 +19,7 @@
 import pytest
 import logging
 import uuid
+import json
 from sagemaker.lineage import (
     action,
     context,
@@ -891,3 +892,17 @@ def _deploy_static_endpoint(execution_arn, sagemaker_session):
             pass
         else:
             raise (e)
+
+
+@pytest.fixture
+def extract_data_from_html():
+    def _method(data):
+        start = data.find("[")
+        end = data.find("]")
+        res = data[start + 1 : end].split("}, ")
+        res = [i + "}" for i in res]
+        res[-1] = res[-1][:-1]
+        data_dict = [json.loads(i) for i in res]
+        return data_dict
+
+    return _method
diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py
index 0c40bbac91..609ba9836d 100644
--- a/tests/integ/sagemaker/lineage/helpers.py
+++ b/tests/integ/sagemaker/lineage/helpers.py
@@ -145,27 +145,23 @@ def clean_all(self):
         for source, dest in self.associations:
             try:
                 self.client.delete_association(SourceArn=source, DestinationArn=dest)
-                time.sleep(0.5)
             except Exception as e:
                 print("skipped " + str(e))
 
         for artifact_arn in self.artifacts:
             try:
                 self.client.delete_artifact(ArtifactArn=artifact_arn)
-                time.sleep(0.5)
             except Exception as e:
                 print("skipped " + str(e))
 
         for action_arn in self.actions:
             try:
                 self.client.delete_action(ActionArn=action_arn)
-                time.sleep(0.5)
             except Exception as e:
                 print("skipped " + str(e))
 
         for context_arn in self.contexts:
             try:
                 self.client.delete_context(ContextArn=context_arn)
-                time.sleep(0.5)
             except Exception as e:
                 print("skipped " + str(e))
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index 1ba87d6a3c..55cddab97b 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -13,7 +13,6 @@
 """This module contains code to test SageMaker ``LineageQueryResult.visualize()``"""
 from __future__ import absolute_import
 import time
-import json
 import os
 
 import pytest
@@ -47,26 +46,22 @@ def test_wide_graph_visualize(sagemaker_session):
     #         \ \--> Artifact
     #          \--->  ...
     try:
-        for i in range(10):
+        for i in range(200):
             artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
             lineage_resource_helper.create_association(
                 source_arn=wide_graph_root_arn, dest_arn=artifact_arn
             )
-    except Exception as e:
-        print(e)
-        lineage_resource_helper.clean_all()
-        assert False
 
-    try:
         lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
         lq_result = lq.query(start_arns=[wide_graph_root_arn])
         lq_result.visualize(path="wideGraph.html")
+
     except Exception as e:
         print(e)
-        lineage_resource_helper.clean_all()
         assert False
 
-    lineage_resource_helper.clean_all()
+    finally:
+        lineage_resource_helper.clean_all()
 
 
 @pytest.mark.skip("visualizer load test")
@@ -84,27 +79,23 @@ def test_long_graph_visualize(sagemaker_session):
                 source_arn=last_arn, dest_arn=new_artifact_arn
             )
             last_arn = new_artifact_arn
-    except Exception as e:
-        print(e)
-        lineage_resource_helper.clean_all()
-        assert False
 
-    try:
         lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
         lq_result = lq.query(
             start_arns=[long_graph_root_arn], direction=LineageQueryDirectionEnum.DESCENDANTS
         )
         # max depth = 10 -> graph rendered only has length of ten (in DESCENDANTS direction)
         lq_result.visualize(path="longGraph.html")
+
     except Exception as e:
         print(e)
-        lineage_resource_helper.clean_all()
         assert False
 
-    lineage_resource_helper.clean_all()
+    finally:
+        lineage_resource_helper.clean_all()
 
 
-def test_graph_visualize(sagemaker_session):
+def test_graph_visualize(sagemaker_session, extract_data_from_html):
     lineage_resource_helper = LineageResourceHelper(sagemaker_session=sagemaker_session)
 
     # create lineage data
@@ -141,24 +132,14 @@ def test_graph_visualize(sagemaker_session):
             dest_arn=endpoint_context,
             association_type="AssociatedWith",
         )
-        time.sleep(1)
-    except Exception as e:
-        print(e)
-        lineage_resource_helper.clean_all()
-        assert False
+        time.sleep(3)
 
-    # visualize
-    try:
+        # visualize
         lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
         lq_result = lq.query(start_arns=[graph_startarn])
         lq_result.visualize(path="testGraph.html")
-    except Exception as e:
-        print(e)
-        lineage_resource_helper.clean_all()
-        assert False
 
-    # check generated graph info
-    try:
+        # check generated graph info
         fo = open("testGraph.html", "r")
         lines = fo.readlines()
         for line in lines:
@@ -167,108 +148,85 @@ def test_graph_visualize(sagemaker_session):
             if "edges = " in line:
                 edge = line
 
-        # extract node data
-        start = node.find("[")
-        end = node.find("]")
-        res = node[start + 1 : end].split("}, ")
-        res = [i + "}" for i in res]
-        res[-1] = res[-1][:-1]
-        node_dict = [json.loads(i) for i in res]
-
-        # extract edge data
-        start = edge.find("[")
-        end = edge.find("]")
-        res = edge[start + 1 : end].split("}, ")
-        res = [i + "}" for i in res]
-        res[-1] = res[-1][:-1]
-        edge_dict = [json.loads(i) for i in res]
+        node_dict = extract_data_from_html(node)
+        edge_dict = extract_data_from_html(edge)
 
         # check node number
         assert len(node_dict) == 5
 
-        # check startarn
-        found_value = next(
-            dictionary for dictionary in node_dict if dictionary["id"] == graph_startarn
-        )
-        assert found_value["color"] == "#146eb4"
-        assert found_value["label"] == "Model"
-        assert found_value["shape"] == "star"
-        assert found_value["title"] == "Artifact"
-
-        # check image artifact
-        found_value = next(
-            dictionary for dictionary in node_dict if dictionary["id"] == image_artifact
-        )
-        assert found_value["color"] == "#146eb4"
-        assert found_value["label"] == "Image"
-        assert found_value["shape"] == "dot"
-        assert found_value["title"] == "Artifact"
-
-        # check dataset artifact
-        found_value = next(
-            dictionary for dictionary in node_dict if dictionary["id"] == dataset_artifact
-        )
-        assert found_value["color"] == "#146eb4"
-        assert found_value["label"] == "DataSet"
-        assert found_value["shape"] == "dot"
-        assert found_value["title"] == "Artifact"
-
-        # check modeldeploy action
-        found_value = next(
-            dictionary for dictionary in node_dict if dictionary["id"] == modeldeploy_action
-        )
-        assert found_value["color"] == "#88c396"
-        assert found_value["label"] == "ModelDeploy"
-        assert found_value["shape"] == "dot"
-        assert found_value["title"] == "Action"
-
-        # check endpoint context
-        found_value = next(
-            dictionary for dictionary in node_dict if dictionary["id"] == endpoint_context
-        )
-        assert found_value["color"] == "#ff9900"
-        assert found_value["label"] == "Endpoint"
-        assert found_value["shape"] == "dot"
-        assert found_value["title"] == "Context"
+        expected_nodes = {
+            graph_startarn: {
+                "color": "#146eb4",
+                "label": "Model",
+                "shape": "star",
+                "title": "Artifact",
+            },
+            image_artifact: {
+                "color": "#146eb4",
+                "label": "Image",
+                "shape": "dot",
+                "title": "Artifact",
+            },
+            dataset_artifact: {
+                "color": "#146eb4",
+                "label": "DataSet",
+                "shape": "dot",
+                "title": "Artifact",
+            },
+            modeldeploy_action: {
+                "color": "#88c396",
+                "label": "ModelDeploy",
+                "shape": "dot",
+                "title": "Action",
+            },
+            endpoint_context: {
+                "color": "#ff9900",
+                "label": "Endpoint",
+                "shape": "dot",
+                "title": "Context",
+            },
+        }
+
+        # check node properties
+        for node in node_dict:
+            for label, val in expected_nodes[node["id"]].items():
+                assert node[label] == val
 
         # check edge number
         assert len(edge_dict) == 4
 
-        # check image_artifact ->  model_artifact(startarn) edge
-        found_value = next(
-            dictionary for dictionary in edge_dict if dictionary["from"] == image_artifact
-        )
-        assert found_value["to"] == graph_startarn
-        assert found_value["title"] == "ContributedTo"
-
-        # check dataset_artifact ->  model_artifact(startarn) edge
-        found_value = next(
-            dictionary for dictionary in edge_dict if dictionary["from"] == dataset_artifact
-        )
-        assert found_value["to"] == graph_startarn
-        assert found_value["title"] == "AssociatedWith"
-
-        # check model_artifact(startarn) ->  modeldeploy_action edge
-        found_value = next(
-            dictionary for dictionary in edge_dict if dictionary["from"] == graph_startarn
-        )
-        assert found_value["to"] == modeldeploy_action
-        assert found_value["title"] == "ContributedTo"
-
-        # check modeldeploy_action ->  endpoint_context edge
-        found_value = next(
-            dictionary for dictionary in edge_dict if dictionary["from"] == modeldeploy_action
-        )
-        assert found_value["to"] == endpoint_context
-        assert found_value["title"] == "AssociatedWith"
+        expected_edges = {
+            image_artifact: {
+                "from": image_artifact,
+                "to": graph_startarn,
+                "title": "ContributedTo",
+            },
+            dataset_artifact: {
+                "from": dataset_artifact,
+                "to": graph_startarn,
+                "title": "AssociatedWith",
+            },
+            graph_startarn: {
+                "from": graph_startarn,
+                "to": modeldeploy_action,
+                "title": "ContributedTo",
+            },
+            modeldeploy_action: {
+                "from": modeldeploy_action,
+                "to": endpoint_context,
+                "title": "AssociatedWith",
+            },
+        }
+
+        # check edge properties
+        for edge in edge_dict:
+            for label, val in expected_edges[edge["from"]].items():
+                assert edge[label] == val
 
     except Exception as e:
         print(e)
-        lineage_resource_helper.clean_all()
-        os.remove("testGraph.html")
         assert False
 
-    # delete generated test graph
-    os.remove("testGraph.html")
-    # clean lineage data
-    lineage_resource_helper.clean_all()
+    finally:
+        lineage_resource_helper.clean_all()
+        os.remove("testGraph.html")

From 23fe1265a8768d03b0bc5c53e879236a3a3d484d Mon Sep 17 00:00:00 2001
From: Yi-Ting Lee <yitingle@amazon.com>
Date: Wed, 3 Aug 2022 15:01:57 -0700
Subject: [PATCH 11/11] sleep time before clean_all added (avoid race
 condition)

---
 tests/integ/sagemaker/lineage/helpers.py                | 4 ++++
 tests/integ/sagemaker/lineage/test_lineage_visualize.py | 9 +++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/tests/integ/sagemaker/lineage/helpers.py b/tests/integ/sagemaker/lineage/helpers.py
index 609ba9836d..3ab22ce332 100644
--- a/tests/integ/sagemaker/lineage/helpers.py
+++ b/tests/integ/sagemaker/lineage/helpers.py
@@ -142,6 +142,10 @@ def create_association(self, source_arn, dest_arn, association_type="AssociatedW
             return False
 
     def clean_all(self):
+        # clean all lineage data created by LineageResourceHelper
+
+        time.sleep(1)  # avoid GSI race condition between create & delete
+
         for source, dest in self.associations:
             try:
                 self.client.delete_association(SourceArn=source, DestinationArn=dest)
diff --git a/tests/integ/sagemaker/lineage/test_lineage_visualize.py b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
index 55cddab97b..d9b3f879a8 100644
--- a/tests/integ/sagemaker/lineage/test_lineage_visualize.py
+++ b/tests/integ/sagemaker/lineage/test_lineage_visualize.py
@@ -29,10 +29,11 @@ def test_LineageResourceHelper(sagemaker_session):
         art1 = lineage_resource_helper.create_artifact(artifact_name=name())
         art2 = lineage_resource_helper.create_artifact(artifact_name=name())
         lineage_resource_helper.create_association(source_arn=art1, dest_arn=art2)
-        lineage_resource_helper.clean_all()
     except Exception as e:
         print(e)
         assert False
+    finally:
+        lineage_resource_helper.clean_all()
 
 
 @pytest.mark.skip("visualizer load test")
@@ -46,7 +47,7 @@ def test_wide_graph_visualize(sagemaker_session):
     #         \ \--> Artifact
     #          \--->  ...
     try:
-        for i in range(200):
+        for i in range(150):
             artifact_arn = lineage_resource_helper.create_artifact(artifact_name=name())
             lineage_resource_helper.create_association(
                 source_arn=wide_graph_root_arn, dest_arn=artifact_arn
@@ -56,6 +57,10 @@ def test_wide_graph_visualize(sagemaker_session):
         lq_result = lq.query(start_arns=[wide_graph_root_arn])
         lq_result.visualize(path="wideGraph.html")
 
+        print("vertex len = ")
+        print(len(lq_result.vertices))
+        assert False
+
     except Exception as e:
         print(e)
         assert False