create context & action added to helper

Author: Yi-Ting Lee

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))

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