code: def training_job_env():
      old_value = os.environ.get(TRAINING_JOB_ARN_ENV)
      os.environ[TRAINING_JOB_ARN_ENV] = "arn:1234aBcDe"
      yield os.environ
      del os.environ[TRAINING_JOB_ARN_ENV]
      if old_value:
          os.environ[TRAINING_JOB_ARN_ENV] = old_value
  
file_path: sagemaker-python-sdk/tests/unit/sagemaker/experiments/test_environment.py
imports: from __future__ import absolute_import
  import json
  import os
  import shutil
  import tempfile
  import unittest.mock
  import pytest
  from sagemaker.experiments import _environment
  from sagemaker.experiments._environment import TRANSFORM_JOB_ARN_ENV, TRAINING_JOB_ARN_ENV
  from sagemaker.utils import retry_with_backoff
label_desc: Use image node_id 2 for calling a global function with example usage: training_job_env() without return types