Source code for ts_lib_pytest.plugin

import pytest

from ts_lib_pytest.df_snapshot import DataframeSnapshot

try:
    from task_script_utils.random import TaskScriptUUIDGenerator
except ImportError:
    pass




def get_uuid(count: int):
    """Create UUID-like string.

    Example: get_uuid(123) -> "00000000-0000-0000-0000-000000000123"
    """
    return f"00000000-0000-0000-0000-{count:0>12d}"





@pytest.fixture
def mock_uuid_generator(monkeypatch):
    """
    Pytest fixture that mocks UUID generation for consistent testing.

    This fixture patches TaskScriptUUIDGenerator and related UUID functions to return
    predictable UUID strings in the format "00000000-0000-0000-0000-00000000000X"
    where X increments by 1 for each call starting from 1.

    The fixture handles ImportError gracefully if task_script_utils is not available.

    Args:
        monkeypatch: Pytest's monkeypatch fixture for patching attributes
    """

    def uuid_gen():
        count = 1
        while True:
            yield get_uuid(count)
            count += 1

    uuid = uuid_gen()
    try:
        TaskScriptUUIDGenerator("Unit Tests", "Unit Tests")
    except NameError:
        pass

    attributes_to_patch = (
        "task_script_utils.random.TSRandom.uuid",
        "task_script_utils.uuid._uuid",
    )
    for attribute in attributes_to_patch:
        try:
            monkeypatch.setattr(attribute, uuid.__next__)
        except (ModuleNotFoundError, AttributeError):
            pass





def pytest_addoption(parser: pytest.Parser):
    """
    Add custom command line options for the ts-lib-pytest plugin.

    This pytest hook adds the --df-snapshot-update option which allows users
    to update DataFrame snapshots instead of comparing against existing ones.

    Args:
        parser: The pytest argument parser to add options to
    """
    group = parser.getgroup("ts-lib-pytest")
    group.addoption(
        "--df-snapshot-update",
        action="store_true",
        help="Update DataFrame snapshots instead of comparing. Implied if using --snapshot-update.",
    )





@pytest.fixture()
def df_snapshot(
    request: pytest.FixtureRequest,
) -> DataframeSnapshot:
    """
    Pytest fixture that provides a DataframeSnapshot instance for snapshot testing.

    This fixture automatically creates a snapshot file path based on the test file
    and test function name, following the pattern:
    `<test_file_stem>_snapshots/<test_name>.parquet`

    The fixture checks for the --df-snapshot-update command line option (or --snapshot-update)
    to determine whether to update snapshots or compare against existing ones.

    To control the snapshot file path, you can set DataframeSnapshot.path accordingly.

    Args:
        request: The pytest FixtureRequest object containing test context

    Returns:
        DataframeSnapshot: An instance configured with the appropriate snapshot file path
            and update mode based on command line options
    """

    # Check if the user has requested to update snapshots
    is_snapshot_update = request.config.getoption("--df-snapshot-update")
    if not is_snapshot_update:
        try:
            # Raises ValueError("no option named '--snapshot-update'") when pytest-snapshot is not installed
            is_snapshot_update = request.config.getoption("--snapshot-update")
        except ValueError:
            pass
    is_snapshot_update = bool(is_snapshot_update)

    # Derive the snapshot path from the test file and test function name
    test_file = request.path
    test_name = request.node.name
    snapshot_dir = test_file.parent.joinpath(test_file.stem + "_snapshots")
    snapshot_path = snapshot_dir.joinpath(test_name + ".parquet")

    return DataframeSnapshot(path=snapshot_path, is_snapshot_update=is_snapshot_update)