piker/tests/conftest.py

from contextlib import asynccontextmanager as acm
from collections.abc import (
    Callable,
    Iterator,
)
from functools import partial
import logging
import os
from pathlib import Path
import sys
from tempfile import TemporaryDirectory
from weakref import (
    ReferenceType,
    ref,
)

# These must be selected before test modules import Qt or Piker.
_original_qt_qpa_platform: str|None = os.environ.get(
    'QT_QPA_PLATFORM',
)
_original_pytest_qt_api: str|None = os.environ.get(
    'PYTEST_QT_API',
)
_headless_qt: bool = '--headless' in sys.argv
if _headless_qt:
    os.environ['QT_QPA_PLATFORM'] = 'offscreen'
else:
    os.environ.setdefault('QT_QPA_PLATFORM', 'offscreen')
os.environ.setdefault('PYTEST_QT_API', 'pyqt6')
_original_xdg_config_home: str|None = os.environ.get(
    'XDG_CONFIG_HOME',
)
_original_xdg_config_dirs: str|None = os.environ.get(
    'XDG_CONFIG_DIRS',
)
_xdg_config_home_owner: TemporaryDirectory[str] = (
    TemporaryDirectory(prefix='piker-pytest-xdg-')
)
_xdg_config_home: Path = Path(_xdg_config_home_owner.name)
_xdg_config_dirs: Path = _xdg_config_home / 'system-config'
_xdg_config_dirs.mkdir()
os.environ['XDG_CONFIG_HOME'] = str(_xdg_config_home)
os.environ['XDG_CONFIG_DIRS'] = str(_xdg_config_dirs)
_xdg_config_restored: bool = False

import pytest
import tractor
from piker import (
    config,
)
from piker.service import (
    Services,
)
from piker.log import get_console_log


# include `tractor`'s built-in fixtures!
pytest_plugins: tuple[str] = (
    "tractor._testing.pytest",
)


@pytest.fixture(scope='session')
def isolated_xdg_config_home() -> Iterator[Path]:
    '''
    Own process-wide config isolation through session teardown.

    '''
    yield _xdg_config_home


def _restore_test_process_environment() -> None:
    '''
    Release process-owned XDG isolation exactly once.

    '''
    global _xdg_config_restored

    if _xdg_config_restored:
        return
    _xdg_config_restored = True

    if _original_xdg_config_home is None:
        os.environ.pop('XDG_CONFIG_HOME', None)
    else:
        os.environ['XDG_CONFIG_HOME'] = _original_xdg_config_home
    if _original_xdg_config_dirs is None:
        os.environ.pop('XDG_CONFIG_DIRS', None)
    else:
        os.environ['XDG_CONFIG_DIRS'] = _original_xdg_config_dirs
    if _original_qt_qpa_platform is None:
        os.environ.pop('QT_QPA_PLATFORM', None)
    else:
        os.environ['QT_QPA_PLATFORM'] = (
            _original_qt_qpa_platform
        )
    if _original_pytest_qt_api is None:
        os.environ.pop('PYTEST_QT_API', None)
    else:
        os.environ['PYTEST_QT_API'] = _original_pytest_qt_api
    _xdg_config_home_owner.cleanup()


def pytest_configure(config: pytest.Config) -> None:
    '''
    Guarantee import-time XDG ownership is released by pytest.

    '''
    config.add_cleanup(_restore_test_process_environment)


@pytest.fixture(scope='session')
def qapp_args(
    isolated_xdg_config_home: Path,

) -> list[str]:
    '''
    Configure pytest-qt after process-wide config isolation.

    '''
    assert isolated_xdg_config_home == _xdg_config_home
    return ['piker-tests']


def pytest_addoption(parser):
    parser.addoption(
        '--headless',
        action='store_true',
        help='Force Qt onto the offscreen platform before imports',
    )
    parser.addoption("--ll", action="store", dest='loglevel',
                     default=None, help="logging level to set when testing")
    parser.addoption("--confdir", default=None,
                     help="Use a practice API account")


@pytest.fixture(scope='session')
def loglevel(request) -> str:
    return request.config.option.loglevel


@pytest.fixture(scope='session')
def test_config():
    dirname = os.path.dirname
    dirpath = os.path.abspath(
        os.path.join(
            dirname(os.path.realpath(__file__)),
            'data'
        )
    )
    return dirpath


@pytest.fixture(scope='session', autouse=True)
def confdir(
    request,
    test_config: str,
):
    '''
    If the `--confdir` flag is not passed use the
    broker config file found in that dir.

    '''
    confdir = request.config.option.confdir
    if confdir is not None:
        config._override_config_dir(confdir)

    return confdir


_ci_env: bool = os.environ.get('CI', False)


@pytest.fixture(scope='session')
def ci_env() -> bool:
    '''
    Detect CI envoirment.

    '''
    return _ci_env


@pytest.fixture()
def log(
    request: pytest.FixtureRequest,
    loglevel: str,
) -> logging.Logger:
    '''
    Deliver a per-test-named ``piker.log`` instance.

    '''
    return get_console_log(
        level=loglevel,
        name=request.node.name,
    )


@pytest.fixture(autouse=True)
def shm_leak_tracker(
    monkeypatch: pytest.MonkeyPatch,

) -> Callable[[], set[str]]:
    '''
    Clean SHM created but not released by the current test process.

    Tractor's token cache includes attachments, so cache membership
    can not prove ownership. Track only successful ``create=True``
    calls through Tractor's own factory. Normal actor-stack teardown
    unlinks these names first; any survivors are exact test-owned leaks
    which this fixture removes before failing the test.

    '''
    from tractor.ipc import _shm

    known_tokens: dict = dict(_shm._known_tokens)
    owned_segments: dict[
        str,
        tuple[
            tuple[int, int]|None,
            ReferenceType,
        ],
    ] = {}
    open_shm = _shm.SharedMemory
    unlink_shm = getattr(_shm, 'shm_unlink', None)

    def normalize_name(name: str) -> str:
        return str(name).lstrip('/')

    def segment_id(segment) -> tuple[int, int]|None:
        fd: int = getattr(segment, '_fd', -1)
        if fd < 0:
            return None
        stat = os.fstat(fd)
        return stat.st_dev, stat.st_ino

    def open_test_shm(*args, **kwargs):
        segment = open_shm(*args, **kwargs)
        create: bool = kwargs.get(
            'create',
            args[1] if len(args) > 1 else False,
        )
        if create:
            name: str = normalize_name(segment.name)
            owned_segments[name] = (
                segment_id(segment),
                ref(segment),
            )
        return segment

    if unlink_shm is not None:
        def unlink_owned_name(name: str) -> None:
            normalized: str = normalize_name(name)
            try:
                unlink_shm(name)
            except FileNotFoundError:
                owned_segments.pop(normalized, None)
                raise
            else:
                owned_segments.pop(normalized, None)

        monkeypatch.setattr(
            _shm,
            'shm_unlink',
            unlink_owned_name,
        )

    def cleanup() -> set[str]:
        leaked: set[str] = set()
        errors: list[Exception] = []
        try:
            for name, record in list(owned_segments.items()):
                expected_id, segment_ref = record
                try:
                    segment = open_shm(
                        name=name,
                        create=False,
                    )
                except FileNotFoundError:
                    segment = None
                except Exception as err:
                    errors.append(err)
                    segment = None
                else:
                    try:
                        actual_id: tuple[int, int]|None = segment_id(
                            segment
                        )
                        if expected_id is None:
                            errors.append(RuntimeError(
                                f'Can not verify SHM ownership: {name}'
                            ))
                        elif actual_id != expected_id:
                            errors.append(RuntimeError(
                                f'SHM name changed ownership: {name}'
                            ))
                        else:
                            leaked.add(name)
                            segment.unlink()
                    except FileNotFoundError:
                        pass
                    except Exception as err:
                        errors.append(err)
                    finally:
                        try:
                            segment.close()
                        except Exception as err:
                            errors.append(err)

                original = segment_ref()
                if original is not None:
                    try:
                        original.close()
                    except Exception as err:
                        errors.append(err)
        finally:
            owned_segments.clear()
            _shm._known_tokens.clear()
            _shm._known_tokens.update(known_tokens)

        if errors:
            raise ExceptionGroup(
                'Failed to clean test-owned SHM',
                errors,
            )
        return leaked

    monkeypatch.setattr(
        _shm,
        'SharedMemory',
        open_test_shm,
    )
    yield cleanup

    leaked: set[str] = cleanup()
    if leaked:
        names: str = '\n'.join(sorted(leaked))
        pytest.fail(
            f'Test leaked SHM segments; cleaned:\n'
            f'{names}'
        )


@acm
async def _open_test_pikerd(
    tmpconfdir: str,
    reg_addr: tuple[str, int|str],
    loglevel: str = 'warning',
    debug_mode: bool = False,

    **kwargs,

) -> tuple[
    str,
    int,
    tractor.Portal
]:
    '''
    Testing helper to startup the service tree and runtime on
    a different port then the default to allow testing alongside
    a running stack.

    Calls `.service._actor_runtime.maybe_open_pikerd()``
    to boot the root actor / tractor runtime.

    '''
    from piker.service import maybe_open_pikerd
    async with (
        maybe_open_pikerd(
            registry_addrs=[reg_addr],
            loglevel=loglevel,

            tractor_runtime_overrides={
                'piker_test_dir': tmpconfdir,
            },

            debug_mode=debug_mode,
            **kwargs,

        ) as service_manager,
    ):
        # this proc/actor is the pikerd
        assert service_manager is Services

        async with tractor.wait_for_actor(
            'pikerd',
            registry_addr=reg_addr,
        ) as portal:
            raddr = portal.chan.raddr
            uw_raddr: tuple = raddr.unwrap()
            assert uw_raddr == reg_addr
            yield (
                raddr._host,
                raddr._port,
                portal,
                service_manager,
            )


@pytest.fixture
def tmpconfdir(
    tmp_path: Path,
) -> Path:
    '''
    Ensure the `brokers.toml` file for the test run exists
    since we changed it to not touch files by default.

    Here we override the default (in the user dir) and
    set the global module var the same as we do inside
    the `tmpconfdir` fixture.

    '''
    tmpconfdir: Path = tmp_path / '_testing'
    tmpconfdir.mkdir()

    # touch the `brokers.toml` file since it won't
    # exist in the tmp test dir by default!
    # override config dir in the root actor (aka
    # this top level testing process).
    from piker import config
    config._config_dir: Path = tmpconfdir

    conf, path = config.load(
        conf_name='brokers',
        touch_if_dne=True,
    )
    assert path.is_file(), 'WTH.. `brokers.toml` not created!?'

    return tmpconfdir

    # NOTE: the `tmp_dir` fixture will wipe any files older then 3 test
    # sessions by default:
    # https://docs.pytest.org/en/6.2.x/tmpdir.html#the-default-base-temporary-directory
    # BUT, if we wanted to always wipe conf dir and all contained files,
    # rmtree(str(tmp_path))


@pytest.fixture
def root_conf(tmpconfdir) -> dict:
    return config.load(
        'conf',
        touch_if_dne=True,
    )


@pytest.fixture
def open_test_pikerd(
    request: pytest.FixtureRequest,
    tmp_path: Path,
    tmpconfdir: Path,

    # XXX from `tractor._testing.pytest` plugin
    loglevel: str,
    reg_addr: tuple,
):

    tmpconfdir_str: str = str(tmpconfdir)

    # NOTE: on linux the tmp config dir is generally located at:
    # /tmp/pytest-of-<username>/pytest-<run#>/test_<current_test_name>/
    # the default `pytest` config ensures that only the last 4 test
    # suite run's dirs will be persisted, otherwise they are removed:
    # https://docs.pytest.org/en/6.2.x/tmpdir.html#the-default-base-temporary-directory
    print(f'CURRENT TEST CONF DIR: {tmpconfdir}')

    conf = request.config
    debug_mode: bool = conf.option.usepdb
    if (
        debug_mode
        and conf.option.capture != 'no'
    ):
        # TODO: how to disable capture dynamically?
        # conf._configured = False
        # conf._do_configure()
        pytest.fail(
            'To use `--pdb` (with `tractor` subactors) you also must also '
            'pass `-s`!'
        )

    yield partial(
        _open_test_pikerd,

        # pass in a unique temp dir for this test request
        # so that we can have multiple tests running (maybe in parallel)
        # bwitout clobbering each other's config state.
        tmpconfdir=tmpconfdir_str,

        # NOTE these come verbatim from `tractor`'s builtin plugin!
        #
        # per-tpt compat registrar address.
        reg_addr=reg_addr,
        # bind in level from fixture.
        # (can be set with `--ll <value>` flag to `pytest`).
        loglevel=loglevel,
        debug_mode=debug_mode,
    )

    # TODO: teardown checks such as,
    # - no leaked subprocs or shm buffers
    # - all requested container service are torn down
    # - certain ``tractor`` runtime state?