plugin.py

import pytest
from pytest_docker_integration.hub import containers_hub, docker_client
from pytest_docker_integration.exceptions import HubStartupFailed


def pytest_addoption(parser):
    group = parser.getgroup("pytest_docker_integration")
    group.addoption(
        "--docker-integration-config",
        action="store",
        nargs="?",
        type=str,
        dest="docker_integration_config",
        help=(
            "Location (as module_path:attr_path string) of the configuration dictionary for the "
            "containers. May point to the dictionary or a callable that returns a dictionary."
        ),
    )
    group.addoption(
        "--docker-log-config-factory",
        action="store",
        nargs="?",
        type=str,
        dest="docker_log_config_factory",
        help=("Location of the callable used for docker LogConfig generation. "),
    )
    group.addoption(
        "--force-pull",
        default=None,
        nargs="?",
        const=True,
        type=str,
        help=(
            "Use '--force-pull' without specifying any parameters if you want all images "
            "to be pulled, even if they exist locally. To pull only selection of the images, use "
            "`--force-pull` with a comma separated service name parameter, "
            "eg.: '--force-pull redis,keycloak'."
        ),
    )


def pytest_configure(config):
    config.addinivalue_line(
        "markers",
        "docker_integration(service_name, setup=None, cleanup=None): "
        "Indicates that the test requires the docker container running for integration tests. "
        "Ensures that the specified docker container is running for tests marked with this mark. "
        "May optionally provide setup and cleanup functions that will run before and after the "
        "marked test (they override the default from the config for a test).",
    )
    containers_hub.configure(config)
    for name in containers_hub.config:
        config.addinivalue_line("markers", f"docker_integration_{name}")


@pytest.hookimpl(hookwrapper=True)
def pytest_runtestloop(session):
    """
    Starts the containers before the tests and removes them after.
    """
    try:
        if not session.config.getoption("collectonly") and not session.testsfailed:
            containers_hub.run_for_items(session.items)
    except HubStartupFailed:
        return pytest.exit(
            "Some of the docker containers failed to start!",
            pytest.ExitCode.INTERNAL_ERROR,
        )
    yield
    containers_hub.remove()
    docker_client.close()


@pytest.hookimpl
def pytest_itemcollected(item):
    """
    Added to allow running only integration tests using specific service.
    Used with -m docker_integration_service_name
    ex. docker_integration_mongo
    """
    for mark in item.iter_markers(name="docker_integration"):
        item.add_marker(f"docker_integration_{mark.args[0]}")


def run_integration_functions(item, func_name):
    """
    Runs function specified in config or marks for all of the docker containers.
    """
    for mark in item.iter_markers(name="docker_integration"):
        # Skip disabled services.
        if mark.args[0] not in containers_hub.managers:
            continue
        setup_func = mark.kwargs.get(func_name, containers_hub.config[mark.args[0]].get(func_name))
        if setup_func:
            setup_func()


@pytest.hookimpl(trylast=True)
def pytest_runtest_setup(item):
    """
    Runs setup function specified in config for all of the docker containers.
    """
    run_integration_functions(item, "setup")


@pytest.hookimpl(trylast=True)
def pytest_runtest_teardown(item, nextitem):
    """
    Runs teardown function specified in config for all of the docker containers.
    """
    if nextitem is None:
        return
    run_integration_functions(item, "cleanup")