You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

95 lines
3.1 KiB

# Ultralytics YOLO 🚀, AGPL-3.0 license
import shutil
from pathlib import Path
import pytest
TMP = Path(__file__).resolve().parent / "tmp" # temp directory for test files
def pytest_addoption(parser):
"""
Add custom command-line options to pytest.
Args:
parser (pytest.config.Parser): The pytest parser object.
"""
parser.addoption("--slow", action="store_true", default=False, help="Run slow tests")
def pytest_configure(config):
"""
Register custom markers to avoid pytest warnings.
Args:
config (pytest.config.Config): The pytest config object.
"""
config.addinivalue_line("markers", "slow: mark test as slow to run")
def pytest_runtest_setup(item):
"""
Setup hook to skip tests marked as slow if the --slow option is not provided.
Args:
item (pytest.Item): The test item object.
"""
if "slow" in item.keywords and not item.config.getoption("--slow"):
pytest.skip("skip slow tests unless --slow is set")
def pytest_collection_modifyitems(config, items):
"""
Modify the list of test items to remove tests marked as slow if the --slow option is not provided.
Args:
config (pytest.config.Config): The pytest config object.
items (list): List of test items to be executed.
"""
if not config.getoption("--slow"):
# Remove the item entirely from the list of test items if it's marked as 'slow'
items[:] = [item for item in items if "slow" not in item.keywords]
def pytest_sessionstart(session):
"""
Initialize session configurations for pytest.
This function is automatically called by pytest after the 'Session' object has been created but before performing
test collection. It sets the initial seeds and prepares the temporary directory for the test session.
Args:
session (pytest.Session): The pytest session object.
"""
from ultralytics.utils.torch_utils import init_seeds
init_seeds()
shutil.rmtree(TMP, ignore_errors=True) # delete any existing tests/tmp directory
TMP.mkdir(parents=True, exist_ok=True) # create a new empty directory
def pytest_terminal_summary(terminalreporter, exitstatus, config):
"""
Cleanup operations after pytest session.
This function is automatically called by pytest at the end of the entire test session. It removes certain files
and directories used during testing.
Args:
terminalreporter (pytest.terminal.TerminalReporter): The terminal reporter object.
exitstatus (int): The exit status of the test run.
config (pytest.config.Config): The pytest config object.
"""
from ultralytics.utils import WEIGHTS_DIR
# Remove files
models = [path for x in ["*.onnx", "*.torchscript"] for path in WEIGHTS_DIR.rglob(x)]
for file in ["bus.jpg", "yolov8n.onnx", "yolov8n.torchscript"] + models:
Path(file).unlink(missing_ok=True)
# Remove directories
models = [path for x in ["*.mlpackage", "*_openvino_model"] for path in WEIGHTS_DIR.rglob(x)]
for directory in [TMP.parents[1] / ".pytest_cache", TMP] + models:
shutil.rmtree(directory, ignore_errors=True)