code: 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",
              },
              {
                  "Arn": "arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext",
                  "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"
          ]
      )
  
      elements = query_response._get_visualization_elements()
  
      assert elements["nodes"][0] == (
          "arn1",
          "Endpoint",
          "Artifact",
          False,
      )
      assert elements["nodes"][1] == ("arn2", "Model", "Context", False)
      assert elements["nodes"][2] == (
          "arn:aws:sagemaker:us-west-2:0123456789012:context/mycontext",
          "Model",
          "Context",
          True,
      )
      assert elements["edges"][0] == ("arn1", "arn2", "Produced")
  
imports: from __future__ import absolute_import
  import unittest.mock
  from sagemaker.lineage.artifact import DatasetArtifact, ModelArtifact, Artifact
  from sagemaker.lineage.context import EndpointContext, Context
  from sagemaker.lineage.action import Action
  from sagemaker.lineage.lineage_trial_component import LineageTrialComponent
  from sagemaker.lineage.query import LineageEntityEnum, LineageSourceEnum, Vertex, LineageQuery
  import pytest
  import re
label_desc: Use image node_id 15 for calling a global function with example usage: test_get_visualization_elements(sagemaker_session) without return types