code: def test_models(sagemaker_session):
      trial_component_arn = "arn:aws:sagemaker:us-west-2:123456789012:trial_component/lineage-unit-3b05f017-0d87-4c37"
      model_arn = (
          "arn:aws:sagemaker:us-west-2:123456789012:context/models"
      )
      model_name = "myDataset"
  
      obj = lineage_trial_component.LineageTrialComponent(
          sagemaker_session,
          trial_component_name="foo",
          trial_component_arn=trial_component_arn,
      )
  
      sagemaker_session.sagemaker_client.query_lineage.return_value = {
          "Vertices": [
              {
                  "Arn": model_arn,
                  "Type": "Model",
                  "LineageType": "Artifact",
              },
          ],
          "Edges": [
              {
                  "SourceArn": "arn1",
                  "DestinationArn": "arn2",
                  "AssociationType": "Produced",
              }
          ],
      }
  
      sagemaker_session.sagemaker_client.describe_artifact.return_value = {
          "ArtifactName": model_name,
          "ArtifactArn": model_arn,
      }
  
      model_list = obj.models()
      expected_calls = [
          unittest.mock.call(
              Direction="Descendants",
              Filters={
                  "Types": ["Model"],
                  "LineageTypes": ["Artifact"],
              },
              IncludeEdges=False,
              MaxDepth=10,
              StartArns=[trial_component_arn],
          ),
      ]
      assert (
          expected_calls
          == sagemaker_session.sagemaker_client.query_lineage.mock_calls
      )
      expected_model_list = [
          artifact.DatasetArtifact(
              artifact_name=model_name,
              artifact_arn=model_arn,
          )
      ]
      assert (
          expected_model_list[0].artifact_arn
          == model_list[0].artifact_arn
      )
      assert (
          expected_model_list[0].artifact_name
          == model_list[0].artifact_name
      )
  
file_path: sagemaker-python-sdk/tests/unit/sagemaker/lineage/test_lineage_trial_component.py
imports: from __future__ import absolute_import
  import unittest.mock
  import pytest
  from sagemaker.lineage import artifact, lineage_trial_component
label_desc: Use image node_id 3 for calling a global function with example usage: test_models(sagemaker_session) without return types