tests.conftest

  1"""> Configuration and fixtures for PyDRex tests."""
  2
  3import sys
  4
  5import matplotlib
  6import numpy as np
  7import pytest
  8from _pytest.logging import LoggingPlugin, _LiveLoggingStreamHandler
  9from pydrex import io as _io
 10from pydrex import logger as _log
 11from pydrex import mock as _mock
 12from pydrex import utils as _utils
 13from scipy.spatial.transform import Rotation
 14
 15from tests import test_vortex_2d as _test_vortex_2d
 16
 17_log.quiet_aliens()  # Stop imported modules from spamming the logs.
 18_, HAS_RAY = _utils.import_proc_pool()
 19if HAS_RAY:
 20    import ray
 21
 22
 23# Set up custom pytest CLI arguments.
 24def pytest_addoption(parser):
 25    parser.addoption(
 26        "--outdir",
 27        metavar="DIR",
 28        default=None,
 29        help="output directory in which to store PyDRex figures/logs",
 30    )
 31    parser.addoption(
 32        "--runslow",
 33        action="store_true",
 34        default=False,
 35        help="run slow tests (HPC cluster recommended, large memory requirement)",
 36    )
 37    parser.addoption(
 38        "--runbig",
 39        action="store_true",
 40        default=False,
 41        help="run tests which are fast enough for home PCs but require 16GB RAM",
 42    )
 43    parser.addoption(
 44        "--ncpus",
 45        default=_utils.default_ncpus(),
 46        type=int,
 47        help="number of CPUs to use for tests that support multiprocessing",
 48    )
 49    parser.addoption(
 50        "--fontsize",
 51        default=None,
 52        type=int,
 53        help="set explicit font size for output figures",
 54    )
 55    parser.addoption(
 56        "--markersize",
 57        default=None,
 58        type=int,
 59        help="set explicit marker size for output figures",
 60    )
 61    parser.addoption(
 62        "--linewidth",
 63        default=None,
 64        type=int,
 65        help="set explicit line width for output figures",
 66    )
 67
 68
 69# The default pytest logging plugin always creates its own handlers...
 70class PytestConsoleLogger(LoggingPlugin):
 71    """Pytest plugin that allows linking up a custom console logger."""
 72
 73    name = "pytest-console-logger"
 74
 75    def __init__(self, config, *args, **kwargs):
 76        super().__init__(config, *args, **kwargs)
 77        terminal_reporter = config.pluginmanager.get_plugin("terminalreporter")
 78        capture_manager = config.pluginmanager.get_plugin("capturemanager")
 79        handler = _LiveLoggingStreamHandler(terminal_reporter, capture_manager)
 80        handler.setFormatter(_log.CONSOLE_LOGGER.formatter)
 81        handler.setLevel(_log.CONSOLE_LOGGER.level)
 82        self.log_cli_handler = handler
 83
 84    # Override original, which tries to delete some silly globals that we aren't
 85    # using anymore, this might break the (already quite broken) -s/--capture.
 86    @pytest.hookimpl(hookwrapper=True)
 87    def pytest_runtest_teardown(self, item):
 88        self.log_cli_handler.set_when("teardown")
 89        yield from self._runtest_for(item, "teardown")
 90
 91
 92@pytest.hookimpl(trylast=True)
 93def pytest_configure(config):
 94    config.addinivalue_line("markers", "slow: mark test as slow to run")
 95    config.addinivalue_line("markers", "big: mark test as requiring 16GB RAM")
 96
 97    # Set custom Matplotlib parameters.
 98    # Alternatively inject a call to `matplotlib.style.use` before starting pytest.
 99    if config.option.fontsize is not None:
100        matplotlib.rcParams["font.size"] = config.option.fontsize
101    if config.option.markersize is not None:
102        matplotlib.rcParams["lines.markersize"] = config.option.markersize
103    if config.option.linewidth is not None:
104        matplotlib.rcParams["lines.linewidth"] = config.option.linewidth
105
106    # Hook up our logging plugin last,
107    # it relies on terminalreporter and capturemanager.
108    if config.option.verbose > 0:
109        terminal_reporter = config.pluginmanager.get_plugin("terminalreporter")
110        capture_manager = config.pluginmanager.get_plugin("capturemanager")
111        handler = _LiveLoggingStreamHandler(terminal_reporter, capture_manager)
112        handler.setFormatter(_log.CONSOLE_LOGGER.formatter)
113        handler.setLevel(_log.CONSOLE_LOGGER.level)
114        _log.LOGGER_PYTEST = handler
115        config.pluginmanager.register(
116            PytestConsoleLogger(config), PytestConsoleLogger.name
117        )
118
119
120def pytest_collection_modifyitems(config, items):
121    if config.getoption("--runslow"):
122        # Don't skip slow tests.
123        _log.info("running slow tests with %d CPUs", config.getoption("--ncpus"))
124    else:
125        skip_slow = pytest.mark.skip(reason="need --runslow option to run")
126        for item in items:
127            if "slow" in item.keywords:
128                item.add_marker(skip_slow)
129
130    if config.getoption("--runbig"):
131        pass  # Don't skip big tests.
132    else:
133        skip_big = pytest.mark.skip(reason="need --runbig option to run")
134        for item in items:
135            if "big" in item.keywords:
136                item.add_marker(skip_big)
137
138
139@pytest.fixture(scope="session")
140def verbose(request):
141    return request.config.option.verbose
142
143
144@pytest.fixture(scope="session")
145def outdir(request):
146    _outdir = request.config.getoption("--outdir")
147    yield _outdir
148    #  Create combined ensemble figure for 2D cell tests after they have all finished.
149    _test_vortex_2d.TestCellOlivineA._make_ensemble_figure(_outdir)
150
151
152@pytest.fixture(scope="session")
153def ncpus(request):
154    return max(1, request.config.getoption("--ncpus"))
155
156
157@pytest.fixture(scope="session")
158def named_tempfile_kwargs(request):
159    if sys.platform == "win32":
160        return {"delete": False}
161    else:
162        return {}
163
164
165@pytest.fixture(scope="session")
166def ray_session():
167    if HAS_RAY:
168        # NOTE: Expects a running Ray cluster with a number of CPUS matching --ncpus.
169        if not ray.is_initialized():
170            ray.init(address="auto")
171            _log.info("using Ray cluster with %s", ray.cluster_resources())
172        yield
173        if ray.is_initialized():
174            ray.shutdown()
175    yield
176
177
178@pytest.fixture(scope="function")
179def console_handler(request):
180    if request.config.option.verbose > 0:  # Show console logs if -v/--verbose given.
181        return request.config.pluginmanager.get_plugin(
182            "pytest-console-logger"
183        ).log_cli_handler
184    return _log.CONSOLE_LOGGER
185
186
187@pytest.fixture
188def params_Fraters2021():
189    return _mock.PARAMS_FRATERS2021
190
191
192@pytest.fixture
193def params_Kaminski2001_fig5_solid():
194    return _mock.PARAMS_KAMINSKI2001_FIG5_SOLID
195
196
197@pytest.fixture
198def params_Kaminski2001_fig5_shortdash():
199    return _mock.PARAMS_KAMINSKI2001_FIG5_SHORTDASH
200
201
202@pytest.fixture
203def params_Kaminski2001_fig5_longdash():
204    return _mock.PARAMS_KAMINSKI2001_FIG5_LONGDASH
205
206
207@pytest.fixture
208def params_Kaminski2004_fig4_triangles():
209    return _mock.PARAMS_KAMINSKI2004_FIG4_TRIANGLES
210
211
212@pytest.fixture
213def params_Kaminski2004_fig4_squares():
214    return _mock.PARAMS_KAMINSKI2004_FIG4_SQUARES
215
216
217@pytest.fixture
218def params_Kaminski2004_fig4_circles():
219    return _mock.PARAMS_KAMINSKI2004_FIG4_CIRCLES
220
221
222@pytest.fixture
223def params_Hedjazian2017():
224    return _mock.PARAMS_HEDJAZIAN2017
225
226
227@pytest.fixture(scope="session")
228def orientations_init_y():
229    rng = np.random.default_rng(seed=8816)
230    return [
231        lambda n_grains: None,  # For random orientations.
232        lambda n_grains: Rotation.from_euler(  # A girdle around Y.
233            "y", [[x * np.pi * 2] for x in rng.random(n_grains)]
234        ).as_matrix(),
235        lambda n_grains: Rotation.from_euler(  # Clustered orientations.
236            "y", [[x * np.pi / 8] for x in rng.random(n_grains)]
237        ).as_matrix(),
238    ]
239
240
241@pytest.fixture(scope="session", params=[100, 500, 1000, 5000, 10000])
242def n_grains(request):
243    return request.param
244
245
246@pytest.fixture(scope="session", params=[[1, 0, 0], [0, 1, 0], [0, 0, 1]])
247def hkl(request):
248    return request.param
249
250
251@pytest.fixture(scope="session", params=["xz", "yz", "xy"])
252def ref_axes(request):
253    return request.param
254
255
256@pytest.fixture(scope="session")
257def seeds():
258    """1000 unique seeds for ensemble runs that need an RNG seed."""
259    return _io.read_scsv(_io.data("rng") / "seeds.scsv").seeds
260
261
262@pytest.fixture(scope="session")
263def seed():
264    """Default seed for test RNG."""
265    return 8816
266
267
268@pytest.fixture(scope="session")
269def seeds_nearX45():
270    """41 seeds which have the initial hexagonal symmetry axis near 45° from X."""
271    return _io.read_scsv(_io.data("rng") / "hexaxis_nearX45_seeds.scsv").seeds