conftest.py

#
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.
#
"""Provides some pytest configuration overrides for both modules."""

# Note: This file is named conftest.py so that pytest picks it up automatically
# without the need to adjust PYTHONPATH or sys.path as much.

import os
import shutil
from tempfile import mkdtemp
from warnings import warn

import pytest
from xdist.workermanage import WorkerController


def is_master(config: pytest.Config) -> bool:
    """True if the code running the given pytest.config object is running in a xdist
    master node or not running xdist at all.
    """
    return not hasattr(config, "workerinput")


def pytest_configure(config: pytest.Config) -> None:
    """Add some additional (global) configuration steps for pytest."""
    # Workaround some issues loading emukit in certain environments.
    if os.environ.get("DISPLAY", None):
        try:
            import matplotlib  # pylint: disable=import-outside-toplevel

            matplotlib.rcParams["backend"] = "agg"
            if is_master(config) or dict(getattr(config, "workerinput", {}))["workerid"] == "gw0":
                # Only warn once.
                warn(
                    UserWarning(
                        "DISPLAY environment variable is set, "
                        "which can cause problems in some setups (e.g. WSL). "
                        "Adjusting matplotlib backend to "
                        f"""'{matplotlib.rcParams["backend"]}' """
                        "to compensate."
                    )
                )
        except ImportError:
            pass

    # Set pandas display options to make inline tests stable.
    import pandas  # pylint: disable=import-outside-toplevel

    pandas.options.display.width = 150
    pandas.options.display.max_columns = 10

    # Create a temporary directory for sharing files between master and worker nodes.
    if is_master(config):
        # Add it to the config so that it can passed to the worker nodes.
        setattr(config, "shared_temp_dir", mkdtemp())


def pytest_configure_node(node: WorkerController) -> None:
    """Xdist hook used to inform workers of the location of the shared temp dir."""
    workerinput: dict = getattr(node, "workerinput")
    workerinput["shared_temp_dir"] = getattr(node.config, "shared_temp_dir")


@pytest.fixture(scope="session")
def shared_temp_dir(request: pytest.FixtureRequest) -> str:
    """Returns a unique and temporary directory which can be shared by master or worker
    nodes in xdist runs.
    """
    if is_master(request.config):
        return str(getattr(request.config, "shared_temp_dir"))
    else:
        workerinput: dict = getattr(request.config, "workerinput")
        return str(workerinput["shared_temp_dir"])


def pytest_unconfigure(config: pytest.Config) -> None:
    """Called after all tests have completed."""
    if is_master(config):
        shared_tmp_dir = getattr(config, "shared_temp_dir", None)
        if shared_tmp_dir:
            shutil.rmtree(str(shared_tmp_dir))


@pytest.fixture(scope="session")
def short_testrun_uid(testrun_uid: str) -> str:
    """Shorten the unique test run id that xdist provides so we can use it with other
    systems (e.g., docker).
    """
    return testrun_uid[0:8]