code: def test_graph_visualize(sagemaker_session, extract_data_from_html):
      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(3)
  
          # visualize
          lq = sagemaker.lineage.query.LineageQuery(sagemaker_session)
          lq_result = lq.query(start_arns=[graph_startarn])
          lq_result.visualize(path="testGraph.html")
  
          # check generated graph info
          fo = open("testGraph.html", "r")
          lines = fo.readlines()
          for line in lines:
              if "nodes = " in line:
                  node = line
              if "edges = " in line:
                  edge = line
  
          node_dict = extract_data_from_html(node)
          edge_dict = extract_data_from_html(edge)
  
          # check node number
          assert len(node_dict) == 5
  
          expected_nodes = {
              graph_startarn: {
                  "color": "#146eb4",
                  "label": "Model",
                  "shape": "star",
                  "title": "Entity: Artifact"
                  + "\nType: Model"
                  + "\nAccount ID: "
                  + str(
                      re.search(r":\d{12}:", graph_startarn).group()[
                          1:-1
                      ]
                  )
                  + "\nName: "
                  + str(re.search(r"\/.*", graph_startarn).group()[1:]),
              },
              image_artifact: {
                  "color": "#146eb4",
                  "label": "Image",
                  "shape": "dot",
                  "title": "Entity: Artifact"
                  + "\nType: Image"
                  + "\nAccount ID: "
                  + str(
                      re.search(r":\d{12}:", image_artifact).group()[
                          1:-1
                      ]
                  )
                  + "\nName: "
                  + str(re.search(r"\/.*", image_artifact).group()[1:]),
              },
              dataset_artifact: {
                  "color": "#146eb4",
                  "label": "Data Set",
                  "shape": "dot",
                  "title": "Entity: Artifact"
                  + "\nType: Data Set"
                  + "\nAccount ID: "
                  + str(
                      re.search(r":\d{12}:", dataset_artifact).group()[
                          1:-1
                      ]
                  )
                  + "\nName: "
                  + str(
                      re.search(r"\/.*", dataset_artifact).group()[1:]
                  ),
              },
              modeldeploy_action: {
                  "color": "#88c396",
                  "label": "Model Deploy",
                  "shape": "dot",
                  "title": "Entity: Action"
                  + "\nType: Model Deploy"
                  + "\nAccount ID: "
                  + str(
                      re.search(
                          r":\d{12}:", modeldeploy_action
                      ).group()[1:-1]
                  )
                  + "\nName: "
                  + str(
                      re.search(r"\/.*", modeldeploy_action).group()[1:]
                  ),
              },
              endpoint_context: {
                  "color": "#ff9900",
                  "label": "Endpoint",
                  "shape": "dot",
                  "title": "Entity: Context"
                  + "\nType: Endpoint"
                  + "\nAccount ID: "
                  + str(
                      re.search(r":\d{12}:", endpoint_context).group()[
                          1:-1
                      ]
                  )
                  + "\nName: "
                  + str(
                      re.search(r"\/.*", endpoint_context).group()[1:]
                  ),
              },
          }
  
          # 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
  
          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)
          assert False
  
      finally:
          lineage_resource_helper.clean_all()
          os.remove("testGraph.html")
  
file_path: sagemaker-python-sdk/tests/integ/sagemaker/lineage/test_lineage_visualize.py
imports: from __future__ import absolute_import
  import time
  import os
  import re
  import pytest
  import sagemaker.lineage.query
  from sagemaker.lineage.query import LineageQueryDirectionEnum
  from tests.integ.sagemaker.lineage.helpers import name, LineageResourceHelper
label_desc: Use image node_id 4 for calling a global function with example usage: test_graph_visualize(sagemaker_session, extract_data_from_html) without return types