# Copyright 2020 The HuggingFace Team. All rights reserved.

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

import collections

import contextlib

import inspect

import logging

import os

import re

import shlex

import shutil

import subprocess

import sys

import tempfile

import unittest

from collections.abc import Mapping

from distutils.util import strtobool

from io import StringIO

from pathlib import Path

from typing import Iterator, List, Union

from unittest import mock

import huggingface_hub

from transformers import logging as transformers_logging

from .deepspeed import is_deepspeed_available

from .integrations import (

is_fairscale_available,

is_optuna_available,

is_ray_available,

is_sigopt_available,

is_wandb_available,

)

from .utils import (

is_accelerate_available,

is_apex_available,

is_bitsandbytes_available,

is_detectron2_available,

is_faiss_available,

is_flax_available,

is_ftfy_available,

is_ipex_available,

is_librosa_available,

is_onnx_available,

is_pandas_available,

is_phonemizer_available,

is_pyctcdecode_available,

is_pytesseract_available,

is_pytorch_quantization_available,

is_rjieba_available,

is_scatter_available,

is_scipy_available,

is_sentencepiece_available,

is_soundfile_availble,

is_spacy_available,

is_tensorflow_probability_available,

is_tensorflow_text_available,

is_tf2onnx_available,

is_tf_available,

is_timm_available,

is_tokenizers_available,

is_torch_available,

is_torch_bf16_cpu_available,

is_torch_bf16_gpu_available,

is_torch_tensorrt_fx_available,

is_torch_tf32_available,

is_torch_tpu_available,

is_torchaudio_available,

is_torchdynamo_available,

is_vision_available,

)

SMALL_MODEL_IDENTIFIER = "julien-c/bert-xsmall-dummy"

DUMMY_UNKNOWN_IDENTIFIER = "julien-c/dummy-unknown"

DUMMY_DIFF_TOKENIZER_IDENTIFIER = "julien-c/dummy-diff-tokenizer"

# Used to test Auto{Config, Model, Tokenizer} model_type detection.

# Used to test the hub

USER = "__DUMMY_TRANSFORMERS_USER__"

ENDPOINT_STAGING = "https://hub-ci.huggingface.co"

# Not critical, only usable on the sandboxed CI instance.

TOKEN = "hf_94wBhPGp6KrrTH3KDchhKpRxZwd6dmHWLL"

def parse_flag_from_env(key, default=False):

try:

value = os.environ[key]

except KeyError:

# KEY isn't set, default to `default`.

_value = default

else:

# KEY is set, convert it to True or False.

try:

_value = strtobool(value)

except ValueError:

# More values are supported, but let's keep the message simple.

raise ValueError(f"If set, {key} must be yes or no.")

return _value

def parse_int_from_env(key, default=None):

try:

value = os.environ[key]

except KeyError:

_value = default

else:

try:

_value = int(value)

except ValueError:

raise ValueError(f"If set, {key} must be a int.")

return _value

_run_slow_tests = parse_flag_from_env("RUN_SLOW", default=False)

_run_pt_tf_cross_tests = parse_flag_from_env("RUN_PT_TF_CROSS_TESTS", default=False)

_run_pt_flax_cross_tests = parse_flag_from_env("RUN_PT_FLAX_CROSS_TESTS", default=False)

_run_custom_tokenizers = parse_flag_from_env("RUN_CUSTOM_TOKENIZERS", default=False)

_run_staging = parse_flag_from_env("HUGGINGFACE_CO_STAGING", default=False)

_run_pipeline_tests = parse_flag_from_env("RUN_PIPELINE_TESTS", default=False)

_run_git_lfs_tests = parse_flag_from_env("RUN_GIT_LFS_TESTS", default=False)

_tf_gpu_memory_limit = parse_int_from_env("TF_GPU_MEMORY_LIMIT", default=None)

def is_pt_tf_cross_test(test_case):

"""

Decorator marking a test as a test that control interactions between PyTorch and TensorFlow.

PT+TF tests are skipped by default and we can run only them by setting RUN_PT_TF_CROSS_TESTS environment variable

to a truthy value and selecting the is_pt_tf_cross_test pytest mark.

"""

if not _run_pt_tf_cross_tests or not is_torch_available() or not is_tf_available():

return unittest.skip("test is PT+TF test")(test_case)

else:

try:

import pytest # We don't need a hard dependency on pytest in the main library

except ImportError:

return test_case

else:

return pytest.mark.is_pt_tf_cross_test()(test_case)

def is_pt_flax_cross_test(test_case):

"""

Decorator marking a test as a test that control interactions between PyTorch and Flax

PT+FLAX tests are skipped by default and we can run only them by setting RUN_PT_FLAX_CROSS_TESTS environment

variable to a truthy value and selecting the is_pt_flax_cross_test pytest mark.

"""

if not _run_pt_flax_cross_tests or not is_torch_available() or not is_flax_available():

return unittest.skip("test is PT+FLAX test")(test_case)

else:

try:

import pytest # We don't need a hard dependency on pytest in the main library

except ImportError:

return test_case

else:

return pytest.mark.is_pt_flax_cross_test()(test_case)

def is_pipeline_test(test_case):

"""

Decorator marking a test as a pipeline test.

Pipeline tests are skipped by default and we can run only them by setting RUN_PIPELINE_TESTS environment variable

to a truthy value and selecting the is_pipeline_test pytest mark.

"""

if not _run_pipeline_tests:

return unittest.skip("test is pipeline test")(test_case)

else:

try:

import pytest # We don't need a hard dependency on pytest in the main library

except ImportError:

return test_case

else:

return pytest.mark.is_pipeline_test()(test_case)

def is_staging_test(test_case):

"""

Decorator marking a test as a staging test.

Those tests will run using the staging environment of huggingface.co instead of the real model hub.

"""

if not _run_staging:

return unittest.skip("test is staging test")(test_case)

else:

try:

import pytest # We don't need a hard dependency on pytest in the main library

except ImportError:

return test_case

else:

return pytest.mark.is_staging_test()(test_case)

def slow(test_case):

"""

Decorator marking a test as slow.

Slow tests are skipped by default. Set the RUN_SLOW environment variable to a truthy value to run them.

"""

return unittest.skipUnless(_run_slow_tests, "test is slow")(test_case)

def tooslow(test_case):

"""

Decorator marking a test as too slow.

Slow tests are skipped while they're in the process of being fixed. No test should stay tagged as "tooslow" as

these will not be tested by the CI.

"""

return unittest.skip("test is too slow")(test_case)

def custom_tokenizers(test_case):

"""

Decorator marking a test for a custom tokenizer.

Custom tokenizers require additional dependencies, and are skipped by default. Set the RUN_CUSTOM_TOKENIZERS

environment variable to a truthy value to run them.

"""

return unittest.skipUnless(_run_custom_tokenizers, "test of custom tokenizers")(test_case)

def require_git_lfs(test_case):

"""

Decorator marking a test that requires git-lfs.

git-lfs requires additional dependencies, and tests are skipped by default. Set the RUN_GIT_LFS_TESTS environment

variable to a truthy value to run them.

"""

return unittest.skipUnless(_run_git_lfs_tests, "test of git lfs workflow")(test_case)

def require_accelerate(test_case):

"""

Decorator marking a test that requires accelerate. These tests are skipped when accelerate isn't installed.

"""

return unittest.skipUnless(is_accelerate_available(), "test requires accelerate")(test_case)

def require_rjieba(test_case):

"""

Decorator marking a test that requires rjieba. These tests are skipped when rjieba isn't installed.

"""

return unittest.skipUnless(is_rjieba_available(), "test requires rjieba")(test_case)

def require_tf2onnx(test_case):

return unittest.skipUnless(is_tf2onnx_available(), "test requires tf2onnx")(test_case)

def require_onnx(test_case):

return unittest.skipUnless(is_onnx_available(), "test requires ONNX")(test_case)

def require_timm(test_case):

"""

Decorator marking a test that requires Timm.

These tests are skipped when Timm isn't installed.

"""

return unittest.skipUnless(is_timm_available(), "test requires Timm")(test_case)

def require_torch(test_case):

"""

Decorator marking a test that requires PyTorch.

These tests are skipped when PyTorch isn't installed.

"""

return unittest.skipUnless(is_torch_available(), "test requires PyTorch")(test_case)

def require_intel_extension_for_pytorch(test_case):

"""

Decorator marking a test that requires Intel Extension for PyTorch.

These tests are skipped when Intel Extension for PyTorch isn't installed or it does not match current PyTorch

version.

"""

return unittest.skipUnless(

is_ipex_available(),

"test requires Intel Extension for PyTorch to be installed and match current PyTorch version, see"

" https://github.com/intel/intel-extension-for-pytorch",

)(test_case)

def require_torch_scatter(test_case):

"""

Decorator marking a test that requires PyTorch scatter.

These tests are skipped when PyTorch scatter isn't installed.

"""

return unittest.skipUnless(is_scatter_available(), "test requires PyTorch scatter")(test_case)

def require_tensorflow_probability(test_case):

"""

Decorator marking a test that requires TensorFlow probability.

These tests are skipped when TensorFlow probability isn't installed.

"""

return unittest.skipUnless(is_tensorflow_probability_available(), "test requires TensorFlow probability")(

test_case

)

def require_torchaudio(test_case):

"""

Decorator marking a test that requires torchaudio. These tests are skipped when torchaudio isn't installed.

"""

return unittest.skipUnless(is_torchaudio_available(), "test requires torchaudio")(test_case)

def require_tf(test_case):

"""

Decorator marking a test that requires TensorFlow. These tests are skipped when TensorFlow isn't installed.

"""

return unittest.skipUnless(is_tf_available(), "test requires TensorFlow")(test_case)

def require_flax(test_case):

"""

Decorator marking a test that requires JAX & Flax. These tests are skipped when one / both are not installed

"""

return unittest.skipUnless(is_flax_available(), "test requires JAX & Flax")(test_case)

def require_sentencepiece(test_case):

"""

Decorator marking a test that requires SentencePiece. These tests are skipped when SentencePiece isn't installed.

"""

return unittest.skipUnless(is_sentencepiece_available(), "test requires SentencePiece")(test_case)

def require_scipy(test_case):

"""

Decorator marking a test that requires Scipy. These tests are skipped when SentencePiece isn't installed.

"""

return unittest.skipUnless(is_scipy_available(), "test requires Scipy")(test_case)

def require_tokenizers(test_case):

"""

Decorator marking a test that requires 🤗 Tokenizers. These tests are skipped when 🤗 Tokenizers isn't installed.

"""

return unittest.skipUnless(is_tokenizers_available(), "test requires tokenizers")(test_case)

def require_tensorflow_text(test_case):

"""

Decorator marking a test that requires tensorflow_text. These tests are skipped when tensroflow_text isn't

installed.

"""

return unittest.skipUnless(is_tensorflow_text_available(), "test requires tensorflow_text")(test_case)

def require_pandas(test_case):

"""

Decorator marking a test that requires pandas. These tests are skipped when pandas isn't installed.

"""

return unittest.skipUnless(is_pandas_available(), "test requires pandas")(test_case)

def require_pytesseract(test_case):

"""

Decorator marking a test that requires PyTesseract. These tests are skipped when PyTesseract isn't installed.

"""

return unittest.skipUnless(is_pytesseract_available(), "test requires PyTesseract")(test_case)

def require_scatter(test_case):

"""

Decorator marking a test that requires PyTorch Scatter. These tests are skipped when PyTorch Scatter isn't

installed.

"""

return unittest.skipUnless(is_scatter_available(), "test requires PyTorch Scatter")(test_case)

def require_pytorch_quantization(test_case):

"""

Decorator marking a test that requires PyTorch Quantization Toolkit. These tests are skipped when PyTorch

Quantization Toolkit isn't installed.

"""

return unittest.skipUnless(is_pytorch_quantization_available(), "test requires PyTorch Quantization Toolkit")(

test_case

)

def require_vision(test_case):

"""

Decorator marking a test that requires the vision dependencies. These tests are skipped when torchaudio isn't

installed.

"""

return unittest.skipUnless(is_vision_available(), "test requires vision")(test_case)

def require_ftfy(test_case):

"""

Decorator marking a test that requires ftfy. These tests are skipped when ftfy isn't installed.

"""

return unittest.skipUnless(is_ftfy_available(), "test requires ftfy")(test_case)

def require_spacy(test_case):

"""

Decorator marking a test that requires SpaCy. These tests are skipped when SpaCy isn't installed.

"""

return unittest.skipUnless(is_spacy_available(), "test requires spacy")(test_case)

def require_torch_multi_gpu(test_case):

"""

Decorator marking a test that requires a multi-GPU setup (in PyTorch). These tests are skipped on a machine without

multiple GPUs.

To run *only* the multi_gpu tests, assuming all test names contain multi_gpu: $ pytest -sv ./tests -k "multi_gpu"

"""

if not is_torch_available():

return unittest.skip("test requires PyTorch")(test_case)

import torch

return unittest.skipUnless(torch.cuda.device_count() > 1, "test requires multiple GPUs")(test_case)

def require_torch_non_multi_gpu(test_case):

"""

Decorator marking a test that requires 0 or 1 GPU setup (in PyTorch).

"""

if not is_torch_available():

return unittest.skip("test requires PyTorch")(test_case)

import torch

return unittest.skipUnless(torch.cuda.device_count() < 2, "test requires 0 or 1 GPU")(test_case)

def require_torch_up_to_2_gpus(test_case):

"""

Decorator marking a test that requires 0 or 1 or 2 GPU setup (in PyTorch).

"""

if not is_torch_available():

return unittest.skip("test requires PyTorch")(test_case)

import torch

return unittest.skipUnless(torch.cuda.device_count() < 3, "test requires 0 or 1 or 2 GPUs")(test_case)

def require_torch_tpu(test_case):

"""

Decorator marking a test that requires a TPU (in PyTorch).

"""

return unittest.skipUnless(is_torch_tpu_available(check_device=False), "test requires PyTorch TPU")(test_case)

if is_torch_available():

# Set env var CUDA_VISIBLE_DEVICES="" to force cpu-mode

import torch

torch_device = "cuda" if torch.cuda.is_available() else "cpu"

else:

torch_device = None

if is_tf_available():

import tensorflow as tf

if is_flax_available():

import jax

jax_device = jax.default_backend()

else:

jax_device = None

def require_torchdynamo(test_case):

"""Decorator marking a test that requires TorchDynamo"""

return unittest.skipUnless(is_torchdynamo_available(), "test requires TorchDynamo")(test_case)

def require_torch_tensorrt_fx(test_case):

"""Decorator marking a test that requires Torch-TensorRT FX"""

return unittest.skipUnless(is_torch_tensorrt_fx_available(), "test requires Torch-TensorRT FX")(test_case)

def require_torch_gpu(test_case):

"""Decorator marking a test that requires CUDA and PyTorch."""

return unittest.skipUnless(torch_device == "cuda", "test requires CUDA")(test_case)

def require_torch_bf16_gpu(test_case):

"""Decorator marking a test that requires torch>=1.10, using Ampere GPU or newer arch with cuda>=11.0"""

return unittest.skipUnless(

is_torch_bf16_gpu_available(),

"test requires torch>=1.10, using Ampere GPU or newer arch with cuda>=11.0",

)(test_case)

def require_torch_bf16_cpu(test_case):

"""Decorator marking a test that requires torch>=1.10, using CPU."""

return unittest.skipUnless(

is_torch_bf16_cpu_available(),

"test requires torch>=1.10, using CPU",

)(test_case)

def require_torch_tf32(test_case):

"""Decorator marking a test that requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7."""

return unittest.skipUnless(

is_torch_tf32_available(), "test requires Ampere or a newer GPU arch, cuda>=11 and torch>=1.7"

)(test_case)

def require_detectron2(test_case):

"""Decorator marking a test that requires detectron2."""

return unittest.skipUnless(is_detectron2_available(), "test requires `detectron2`")(test_case)

def require_faiss(test_case):

"""Decorator marking a test that requires faiss."""

return unittest.skipUnless(is_faiss_available(), "test requires `faiss`")(test_case)

def require_optuna(test_case):

"""

Decorator marking a test that requires optuna.

These tests are skipped when optuna isn't installed.

"""

return unittest.skipUnless(is_optuna_available(), "test requires optuna")(test_case)

def require_ray(test_case):

"""

Decorator marking a test that requires Ray/tune.

These tests are skipped when Ray/tune isn't installed.

"""

return unittest.skipUnless(is_ray_available(), "test requires Ray/tune")(test_case)

def require_sigopt(test_case):

"""

Decorator marking a test that requires SigOpt.

These tests are skipped when SigOpt isn't installed.

"""

return unittest.skipUnless(is_sigopt_available(), "test requires SigOpt")(test_case)

def require_wandb(test_case):

"""

Decorator marking a test that requires wandb.

These tests are skipped when wandb isn't installed.

"""

return unittest.skipUnless(is_wandb_available(), "test requires wandb")(test_case)

def require_soundfile(test_case):

"""

Decorator marking a test that requires soundfile

These tests are skipped when soundfile isn't installed.

"""

return unittest.skipUnless(is_soundfile_availble(), "test requires soundfile")(test_case)

def require_deepspeed(test_case):

"""

Decorator marking a test that requires deepspeed

"""

return unittest.skipUnless(is_deepspeed_available(), "test requires deepspeed")(test_case)

def require_fairscale(test_case):

"""

Decorator marking a test that requires fairscale

"""

return unittest.skipUnless(is_fairscale_available(), "test requires fairscale")(test_case)

def require_apex(test_case):

"""

Decorator marking a test that requires apex

"""

return unittest.skipUnless(is_apex_available(), "test requires apex")(test_case)

def require_bitsandbytes(test_case):

"""

Decorator for bits and bytes (bnb) dependency

"""

return unittest.skipUnless(is_bitsandbytes_available(), "test requires bnb")(test_case)

def require_phonemizer(test_case):

"""

Decorator marking a test that requires phonemizer

"""

return unittest.skipUnless(is_phonemizer_available(), "test requires phonemizer")(test_case)

def require_pyctcdecode(test_case):

"""

Decorator marking a test that requires pyctcdecode

"""

return unittest.skipUnless(is_pyctcdecode_available(), "test requires pyctcdecode")(test_case)

def require_librosa(test_case):

"""

Decorator marking a test that requires librosa

"""

return unittest.skipUnless(is_librosa_available(), "test requires librosa")(test_case)

def cmd_exists(cmd):

return shutil.which(cmd) is not None

def require_usr_bin_time(test_case):

"""

Decorator marking a test that requires `/usr/bin/time`

"""

return unittest.skipUnless(cmd_exists("/usr/bin/time"), "test requires /usr/bin/time")(test_case)

def get_gpu_count():

"""

Return the number of available gpus (regardless of whether torch, tf or jax is used)

"""

if is_torch_available():

import torch

return torch.cuda.device_count()

elif is_tf_available():

import tensorflow as tf

return len(tf.config.list_physical_devices("GPU"))

elif is_flax_available():

import jax

return jax.device_count()

else:

return 0

def get_tests_dir(append_path=None):

"""

Args:

append_path: optional path to append to the tests dir path

Return:

The full path to the `tests` dir, so that the tests can be invoked from anywhere. Optionally `append_path` is

joined after the `tests` dir the former is provided.

"""

# this function caller's __file__

caller__file__ = inspect.stack()[1][1]

tests_dir = os.path.abspath(os.path.dirname(caller__file__))

while not tests_dir.endswith("tests"):

tests_dir = os.path.dirname(tests_dir)

if append_path:

return os.path.join(tests_dir, append_path)

else:

return tests_dir

#

# Helper functions for dealing with testing text outputs

# The original code came from:

# https://github.com/fastai/fastai/blob/master/tests/utils/text.py

# When any function contains print() calls that get overwritten, like progress bars,

# a special care needs to be applied, since under pytest -s captured output (capsys

# or contextlib.redirect_stdout) contains any temporary printed strings, followed by

# \r's. This helper function ensures that the buffer will contain the same output

# with and without -s in pytest, by turning:

# foo bar\r tar mar\r final message

# into:

# final message

# it can handle a single string or a multiline buffer

def apply_print_resets(buf):

return re.sub(r"^.*\r", "", buf, 0, re.M)

def assert_screenout(out, what):

out_pr = apply_print_resets(out).lower()

match_str = out_pr.find(what.lower())

assert match_str != -1, f"expecting to find {what} in output: f{out_pr}"

class CaptureStd:

"""

Context manager to capture:

- stdout: replay it, clean it up and make it available via `obj.out`

- stderr: replay it and make it available via `obj.err`

Args:

out (`bool`, *optional*, defaults to `True`): Whether to capture stdout or not.

err (`bool`, *optional*, defaults to `True`): Whether to capture stderr or not.

replay (`bool`, *optional*, defaults to `True`): Whether to replay or not.

By default each captured stream gets replayed back on context's exit, so that one can see what the test was

doing. If this is a not wanted behavior and the captured data shouldn't be replayed, pass `replay=False` to

disable this feature.

Examples:

```python

# to capture stdout only with auto-replay

with CaptureStdout() as cs:

print("Secret message")

assert "message" in cs.out

# to capture stderr only with auto-replay

import sys

with CaptureStderr() as cs:

print("Warning: ", file=sys.stderr)

assert "Warning" in cs.err

# to capture both streams with auto-replay

with CaptureStd() as cs:

print("Secret message")

print("Warning: ", file=sys.stderr)

assert "message" in cs.out

assert "Warning" in cs.err

# to capture just one of the streams, and not the other, with auto-replay

with CaptureStd(err=False) as cs:

print("Secret message")

assert "message" in cs.out

# but best use the stream-specific subclasses

# to capture without auto-replay

with CaptureStd(replay=False) as cs:

print("Secret message")

assert "message" in cs.out

```"""

def __init__(self, out=True, err=True, replay=True):

self.replay = replay

if out:

self.out_buf = StringIO()

self.out = "error: CaptureStd context is unfinished yet, called too early"

else:

self.out_buf = None

self.out = "not capturing stdout"

if err:

self.err_buf = StringIO()

self.err = "error: CaptureStd context is unfinished yet, called too early"

else:

self.err_buf = None

self.err = "not capturing stderr"

def __enter__(self):

if self.out_buf:

self.out_old = sys.stdout

sys.stdout = self.out_buf

if self.err_buf:

self.err_old = sys.stderr

sys.stderr = self.err_buf

return self

def __exit__(self, *exc):

if self.out_buf:

sys.stdout = self.out_old

captured = self.out_buf.getvalue()

if self.replay:

sys.stdout.write(captured)

self.out = apply_print_resets(captured)

if self.err_buf:

sys.stderr = self.err_old

captured = self.err_buf.getvalue()

if self.replay:

sys.stderr.write(captured)

self.err = captured

def __repr__(self):

msg = ""

if self.out_buf:

msg += f"stdout: {self.out}\n"

if self.err_buf:

msg += f"stderr: {self.err}\n"

return msg

# in tests it's the best to capture only the stream that's wanted, otherwise

# it's easy to miss things, so unless you need to capture both streams, use the

# subclasses below (less typing). Or alternatively, configure `CaptureStd` to

# disable the stream you don't need to test.

class CaptureStdout(CaptureStd):

"""Same as CaptureStd but captures only stdout"""

def __init__(self, replay=True):

super().__init__(err=False, replay=replay)

class CaptureStderr(CaptureStd):

"""Same as CaptureStd but captures only stderr"""

def __init__(self, replay=True):

super().__init__(out=False, replay=replay)

class CaptureLogger:

"""

Context manager to capture `logging` streams

Args:

logger: 'logging` logger object

Returns:

The captured output is available via `self.out`

Example:

```python

>>> from transformers import logging

>>> from transformers.testing_utils import CaptureLogger

>>> msg = "Testing 1, 2, 3"

>>> logging.set_verbosity_info()

>>> logger = logging.get_logger("transformers.models.bart.tokenization_bart")

>>> with CaptureLogger(logger) as cl:

... logger.info(msg)

>>> assert cl.out, msg + "\n"

```

"""

def __init__(self, logger):

self.logger = logger

self.io = StringIO()

self.sh = logging.StreamHandler(self.io)

self.out = ""

def __enter__(self):

self.logger.addHandler(self.sh)

return self

def __exit__(self, *exc):

self.logger.removeHandler(self.sh)

self.out = self.io.getvalue()

def __repr__(self):

return f"captured: {self.out}\n"

@contextlib.contextmanager

def LoggingLevel(level):

"""

This is a context manager to temporarily change transformers modules logging level to the desired value and have it

restored to the original setting at the end of the scope.

Example:

```python

with LoggingLevel(logging.INFO):

AutoModel.from_pretrained("gpt2") # calls logger.info() several times

```

"""

orig_level = transformers_logging.get_verbosity()

try:

transformers_logging.set_verbosity(level)

yield

finally:

transformers_logging.set_verbosity(orig_level)

@contextlib.contextmanager

# adapted from https://stackoverflow.com/a/64789046/9201239

def ExtendSysPath(path: Union[str, os.PathLike]) -> Iterator[None]:

"""

Temporary add given path to `sys.path`.

Usage :

```python

with ExtendSysPath("/path/to/dir"):

mymodule = importlib.import_module("mymodule")

```

"""

path = os.fspath(path)

try:

sys.path.insert(0, path)

yield

finally:

sys.path.remove(path)

class TestCasePlus(unittest.TestCase):

"""

This class extends *unittest.TestCase* with additional features.

Feature 1: A set of fully resolved important file and dir path accessors.

In tests often we need to know where things are relative to the current test file, and it's not trivial since the

test could be invoked from more than one directory or could reside in sub-directories with different depths. This

class solves this problem by sorting out all the basic paths and provides easy accessors to them:

- `pathlib` objects (all fully resolved):

- `test_file_path` - the current test file path (=`__file__`)

- `test_file_dir` - the directory containing the current test file

- `tests_dir` - the directory of the `tests` test suite

- `examples_dir` - the directory of the `examples` test suite

- `repo_root_dir` - the directory of the repository

- `src_dir` - the directory of `src` (i.e. where the `transformers` sub-dir resides)

- stringified paths---same as above but these return paths as strings, rather than `pathlib` objects:

- `test_file_path_str`

- `test_file_dir_str`

- `tests_dir_str`

- `examples_dir_str`

- `repo_root_dir_str`

- `src_dir_str`

Feature 2: Flexible auto-removable temporary dirs which are guaranteed to get removed at the end of test.

1. Create a unique temporary dir:

```python

def test_whatever(self):

tmp_dir = self.get_auto_remove_tmp_dir()

```

`tmp_dir` will contain the path to the created temporary dir. It will be automatically removed at the end of the

test.

2. Create a temporary dir of my choice, ensure it's empty before the test starts and don't

empty it after the test.

```python

def test_whatever(self):

tmp_dir = self.get_auto_remove_tmp_dir("./xxx")

```

This is useful for debug when you want to monitor a specific directory and want to make sure the previous tests

didn't leave any data in there.

3. You can override the first two options by directly overriding the `before` and `after` args, leading to the

following behavior:

`before=True`: the temporary dir will always be cleared at the beginning of the test.

`before=False`: if the temporary dir already existed, any existing files will remain there.

`after=True`: the temporary dir will always be deleted at the end of the test.

`after=False`: the temporary dir will always be left intact at the end of the test.

Note 1: In order to run the equivalent of `rm -r` safely, only subdirs of the project repository checkout are

allowed if an explicit `tmp_dir` is used, so that by mistake no `/tmp` or similar important part of the filesystem

will get nuked. i.e. please always pass paths that start with `./`

Note 2: Each test can register multiple temporary dirs and they all will get auto-removed, unless requested