conftest.py

# This file is part of IMAS-Python.
# You should have received the IMAS-Python LICENSE file with this project.
#
# Set up pytest:
# - Backend parametrization (and corresponding command line options)
# - IDS name parametrization (and corresponding command line options)
# - Fixtures that are useful across test modules

import functools
import logging
import os
import sys
from copy import deepcopy
from pathlib import Path

try:
    from importlib.resources import files
except ImportError:  # Python 3.8 support
    from importlib_resources import files

import numpy as np
import pytest
from packaging.version import Version

from imas.backends.imas_core.imas_interface import has_imas as _has_imas
from imas.backends.imas_core.imas_interface import ll_interface, lowlevel
from imas.dd_zip import dd_etree, dd_xml_versions, latest_dd_version
from imas.ids_defs import (
    ASCII_BACKEND,
    HDF5_BACKEND,
    IDS_TIME_MODE_INDEPENDENT,
    MDSPLUS_BACKEND,
    MEMORY_BACKEND,
)
from imas.ids_factory import IDSFactory

logger = logging.getLogger("imas")
logger.setLevel(logging.INFO)

os.environ["IMAS_AL_DISABLE_VALIDATE"] = "1"


try:
    import imas  # noqa
except ImportError:

    class SkipOnIMASAccess:
        def __getattr__(self, attr):
            pytest.skip("This test requires the `imas` HLI, which is not available.")

    # Any test that tries to access an attribute from the `imas` package will be skipped
    sys.modules["imas"] = SkipOnIMASAccess()


def pytest_addoption(parser):
    # if none of these are specified, test with all backends
    parser.addoption("--mdsplus", action="store_true", help="test with MDSPlus backend")
    parser.addoption("--memory", action="store_true", help="test with memory backend")
    parser.addoption("--ascii", action="store_true", help="test with ascii backend")
    parser.addoption("--hdf5", action="store_true", help="test with HDF5 backend")
    parser.addoption("--mini", action="store_true", help="small test with few types")
    parser.addoption(
        "--ids", action="append", help="small test with few types", nargs="+"
    )


_BACKENDS = {
    "ascii": ASCII_BACKEND,
    "memory": MEMORY_BACKEND,
    "hdf5": HDF5_BACKEND,
    "mdsplus": MDSPLUS_BACKEND,
}


try:
    import pytest_xdist
except ImportError:
    # If pytest-xdist is not available we provide a dummy worker_id fixture.
    @pytest.fixture()
    def worker_id():
        return "master"


@pytest.fixture(params=_BACKENDS)
def backend(pytestconfig: pytest.Config, request: pytest.FixtureRequest):
    backends_provided = any(map(pytestconfig.getoption, _BACKENDS))
    if not _has_imas:
        if backends_provided:
            raise RuntimeError(
                "Explicit backends are provided, but IMAS is not available."
            )
        pytest.skip("No IMAS available, skip tests using a backend")
    if backends_provided and not pytestconfig.getoption(request.param):
        pytest.skip(f"Tests for {request.param} backend are skipped.")
    return _BACKENDS[request.param]


@pytest.fixture()
def has_imas():
    return _has_imas


@pytest.fixture()
def requires_imas():
    if not _has_imas:
        pytest.skip("No IMAS available")


def pytest_generate_tests(metafunc):
    if "ids_name" in metafunc.fixturenames:
        if metafunc.config.getoption("ids"):
            ids_names = [
                item
                for arg in metafunc.config.getoption("ids")
                for item in arg[0].split(",")
            ]
            metafunc.parametrize("ids_name", ids_names)
        elif metafunc.config.getoption("mini"):
            metafunc.parametrize("ids_name", ["pulse_schedule"])
        else:
            metafunc.parametrize("ids_name", list(IDSFactory()))


@pytest.fixture()
def latest_factory():
    latest_version = latest_dd_version()
    default_factory = IDSFactory()
    default_version = default_factory.version
    # default_version might be a dev version (generated by `git describe`), which cannot
    # be parsed by Version. Take the part before the '-':
    if "-" in default_version:
        default_version = default_version[: default_version.find("-")]
    if Version(default_version) >= Version(latest_version):
        return default_factory
    return IDSFactory(latest_version)


@pytest.fixture()
def latest_factory3():
    """Get most recent DDv3 version."""
    for version in reversed(dd_xml_versions()):
        if version.startswith("3."):
            return IDSFactory(version)


# Fixtures for various assets
@pytest.fixture()
def imas_assets():
    return files("imas") / "assets"


@pytest.fixture()
def fake_toplevel_xml(imas_assets):
    return imas_assets / "IDS_fake_toplevel.xml"


@pytest.fixture()
def ids_minimal(imas_assets):
    return imas_assets / "IDS_minimal.xml"


@pytest.fixture()
def ids_minimal2(imas_assets):
    return imas_assets / "IDS_minimal_2.xml"


@pytest.fixture()
def ids_minimal_struct_array(imas_assets):
    return imas_assets / "IDS_minimal_struct_array.xml"


@pytest.fixture()
def ids_minimal_types(imas_assets):
    return imas_assets / "IDS_minimal_types.xml"


@pytest.fixture
def fake_structure_xml(fake_toplevel_xml):
    tree = dd_etree(version=None, xml_path=fake_toplevel_xml)
    return deepcopy(tree.find("IDS"))


@pytest.fixture
def fake_filled_toplevel(fake_toplevel_xml: Path, worker_id: str, tmp_path: Path):
    """A very specifically filled smallish toplevel"""
    factory = IDSFactory(xml_path=fake_toplevel_xml)
    top = factory.new("gyrokinetics")
    top.wavevector.resize(1)
    top.wavevector[0].eigenmode.resize(1)
    eig = top.wavevector[0].eigenmode[0]
    eig.frequency_norm = 10
    eig.poloidal_angle = np.linspace(0, 2, num=10) * np.pi
    top.ids_properties.homogeneous_time = IDS_TIME_MODE_INDEPENDENT

    yield top


def _lowlevel_wrapper(original_method):
    @functools.wraps(original_method)
    def wrapper(*args, **kwargs):
        result = original_method(*args, **kwargs)
        name = original_method.__name__
        logger.info("UAL lowlevel call: %r(%s, %s) -> %s", name, args, kwargs, result)
        return result

    return wrapper


@pytest.fixture
def log_lowlevel_calls(monkeypatch, requires_imas):
    """Debugging fixture to log calls to the imas lowlevel module."""
    for al_function in dir(lowlevel):
        if al_function.startswith("ual_") or al_function.startswith("al"):
            wrapper = _lowlevel_wrapper(getattr(lowlevel, al_function))
            monkeypatch.setattr(lowlevel, al_function, wrapper)
    for al_function in dir(ll_interface):
        if not al_function.startswith("_"):
            wrapper = _lowlevel_wrapper(getattr(ll_interface, al_function))
            monkeypatch.setattr(ll_interface, al_function, wrapper)