code: def with_environ(key: str, value: Optional[str]):
      """Runs the given test with envrionment variables set."""
  
      def set_env(new_value):
          if new_value is None:
              if key in os.environ:
                  del os.environ[key]
          else:
              os.environ[key] = new_value
  
      def decorator(f):
          def wrapper(*a, **k):
              value_before = os.environ.get(key, None)
              set_env(value)
              try:
                  return f(*a, **k)
              finally:
                  set_env(value_before)
  
          return wrapper
  
      return decorator
  
file_path: dm-haiku/haiku/_src/test_utils.py
imports: from collections.abc import Generator, Sequence
  import functools
  import inspect
  import itertools
  import os
  import types
  from typing import Any, Callable, Optional, TypeVar
  from absl.testing import parameterized
  from haiku._src import config
  from haiku._src import transform
  import jax
label_desc: Use image node_id 7 for calling a global function with example usage: with_environ(key, value) and returns: decorator, wrapper, f