python_packaging/src/gmxapi/testsupport.py

   1 #
   2 # This file is part of the GROMACS molecular simulation package.
   3 #
   4 # Copyright (c) 2019,2020,2021, by the GROMACS development team, led by
   5 # Mark Abraham, David van der Spoel, Berk Hess, and Erik Lindahl,
   6 # and including many others, as listed in the AUTHORS file in the
   7 # top-level source directory and at http://www.gromacs.org.
   8 #
   9 # GROMACS is free software; you can redistribute it and/or
  10 # modify it under the terms of the GNU Lesser General Public License
  11 # as published by the Free Software Foundation; either version 2.1
  12 # of the License, or (at your option) any later version.
  13 #
  14 # GROMACS is distributed in the hope that it will be useful,
  15 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  16 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  17 # Lesser General Public License for more details.
  18 #
  19 # You should have received a copy of the GNU Lesser General Public
  20 # License along with GROMACS; if not, see
  21 # http://www.gnu.org/licenses, or write to the Free Software Foundation,
  22 # Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA.
  23 #
  24 # If you want to redistribute modifications to GROMACS, please
  25 # consider that scientific software is very special. Version
  26 # control is crucial - bugs must be traceable. We will be happy to
  27 # consider code for inclusion in the official distribution, but
  28 # derived work must not be called official GROMACS. Details are found
  29 # in the README & COPYING files - if they are missing, get the
  30 # official version at http://www.gromacs.org.
  31 #
  32 # To help us fund GROMACS development, we humbly ask that you cite
  33 # the research papers on the package. Check out http://www.gromacs.org.
  34
  35 """Reusable definitions for test modules.
  36
  37 Provides utilities and pytest fixtures for gmxapi and GROMACS tests.
  38
  39 To load these facilities in a pytest environment, set a `pytest_plugins`
  40 variable in a conftest.py
  41 (Reference https://docs.pytest.org/en/latest/writing_plugins.html#requiring-loading-plugins-in-a-test-module-or-conftest-file)
  42
  43     pytest_plugins = "gmxapi.testsupport"
  44
  45 .. seealso:: https://docs.pytest.org/en/latest/plugins.html#findpluginname
  46
  47 .. todo:: Consider moving this to a separate optional package.
  48 """
  49
  50 import json
  51 import logging
  52 import os
  53 import shutil
  54 import tempfile
  55 import warnings
  56 from contextlib import contextmanager
  57 from enum import Enum
  58 from typing import Union
  59
  60 import pytest
  61
  62 mpi_status = 'Test requires mpi4py managing 2 MPI ranks.'
  63 skip_mpi = False
  64 try:
  65     from mpi4py import MPI
  66
  67     if not MPI.Is_initialized():
  68         skip_mpi = True
  69         mpi_status += ' MPI is not initialized'
  70     elif MPI.COMM_WORLD.Get_size() < 2:
  71         skip_mpi = True
  72         mpi_status += ' MPI context is too small.'
  73 except ImportError:
  74     skip_mpi = True
  75     mpi_status += ' mpi4py is not available.'
  76
  77
  78 def pytest_configure(config):
  79     config.addinivalue_line("markers", "withmpi_only: test requires mpi4py managing 2 MPI ranks.")
  80
  81
  82 def pytest_runtest_setup(item):
  83     # Handle the withmpi_only marker.
  84     for _ in item.iter_markers(name='withmpi_only'):
  85         if skip_mpi:
  86             pytest.skip(mpi_status)
  87         # The API uses iteration because markers may be duplicated, but we only
  88         # care about whether 'withmpi_only' occurs at all.
  89         break
  90
  91
  92 def pytest_addoption(parser):
  93     """Add command-line user options for the pytest invocation."""
  94     parser.addoption(
  95         '--rm',
  96         action='store',
  97         default='always',
  98         choices=['always', 'never', 'success'],
  99         help='Remove temporary directories "always", "never", or on "success".'
 100     )
 101     parser.addoption(
 102         '--threads',
 103         type=int,
 104         help='Maximum number of threads per process per gmxapi session.'
 105     )
 106
 107
 108 class RmOption(Enum):
 109     """Enumerate allowable values of the --rm option."""
 110     always = 'always'
 111     never = 'never'
 112     success = 'success'
 113
 114
 115 @pytest.fixture(scope='session')
 116 def remove_tempdir(request) -> RmOption:
 117     """pytest fixture to get access to the --rm CLI option."""
 118     arg = request.config.getoption('--rm')
 119     return RmOption(arg)
 120
 121
 122 @pytest.fixture(scope='session')
 123 def gmxconfig():
 124     from .commandline import _config
 125     config = _config()
 126     yield config
 127
 128
 129 @pytest.fixture(scope='session')
 130 def mdrun_kwargs(request, gmxconfig):
 131     """pytest fixture to provide a mdrun_kwargs dictionary for the mdrun ResourceManager.
 132     """
 133     from gmxapi.simulation.mdrun import ResourceManager as _ResourceManager
 134     if gmxconfig is None:
 135         raise RuntimeError('--threads argument requires a usable gmxconfig.json')
 136     arg = request.config.getoption('--threads')
 137     if arg is None:
 138         return {}
 139     mpi_type = gmxconfig['gmx_mpi_type']
 140     if mpi_type is not None and mpi_type == "tmpi":
 141         kwargs = {'threads': int(arg)}
 142     else:
 143         kwargs = {}
 144     # TODO: (#3718) Normalize the handling of run-time arguments.
 145     _ResourceManager.mdrun_kwargs = dict(**kwargs)
 146     return kwargs
 147
 148
 149 @contextmanager
 150 def scoped_chdir(dir):
 151     oldpath = os.getcwd()
 152     os.chdir(dir)
 153     try:
 154         yield dir
 155         # If the `with` block using scoped_chdir produces an exception, it will
 156         # be raised at this point in this function. We want the exception to
 157         # propagate out of the `with` block, but first we want to restore the
 158         # original working directory, so we skip `except` but provide a `finally`.
 159     finally:
 160         os.chdir(oldpath)
 161
 162
 163 @contextmanager
 164 def _cleandir(remove_tempdir: Union[str, RmOption]):
 165     """Context manager for a clean temporary working directory.
 166
 167     Arguments:
 168         remove_tempdir (RmOption): whether to remove temporary directory "always",
 169                                    "never", or on "success"
 170
 171     Raises:
 172         ValueError: if remove_tempdir value is not valid.
 173
 174     The context manager will issue a warning for each temporary directory that
 175     is not removed.
 176     """
 177     if not isinstance(remove_tempdir, RmOption):
 178         remove_tempdir = RmOption(remove_tempdir)
 179
 180     newpath = tempfile.mkdtemp()
 181
 182     def remove():
 183         shutil.rmtree(newpath)
 184
 185     def warn():
 186         warnings.warn('Temporary directory not removed: {}'.format(newpath))
 187
 188     # Initialize callback function reference
 189     if remove_tempdir == RmOption.always:
 190         callback = remove
 191     else:
 192         callback = warn
 193
 194     try:
 195         with scoped_chdir(newpath):
 196             yield newpath
 197         # If we get to this line, the `with` block using _cleandir did not throw.
 198         # Clean up the temporary directory unless the user specified `--rm never`.
 199         # I.e. If the user specified `--rm success`, then we need to toggle from `warn` to `remove`.
 200         if remove_tempdir != RmOption.never:
 201             callback = remove
 202     finally:
 203         callback()
 204
 205
 206 @pytest.fixture
 207 def cleandir(remove_tempdir: RmOption):
 208     """Provide a clean temporary working directory for a test.
 209
 210     Example usage:
 211
 212         import os
 213         import pytest
 214
 215         @pytest.mark.usefixtures("cleandir")
 216         def test_cwd_starts_empty():
 217             assert os.listdir(os.getcwd()) == []
 218             with open("myfile", "w") as f:
 219                 f.write("hello")
 220
 221         def test_cwd_also_starts_empty(cleandir):
 222             assert os.listdir(os.getcwd()) == []
 223             assert os.path.abspath(os.getcwd()) == os.path.abspath(cleandir)
 224             with open("myfile", "w") as f:
 225                 f.write("hello")
 226
 227         @pytest.mark.usefixtures("cleandir")
 228         class TestDirectoryInit(object):
 229             def test_cwd_starts_empty(self):
 230                 assert os.listdir(os.getcwd()) == []
 231                 with open("myfile", "w") as f:
 232                     f.write("hello")
 233
 234             def test_cwd_also_starts_empty(self):
 235                 assert os.listdir(os.getcwd()) == []
 236                 with open("myfile", "w") as f:
 237                     f.write("hello")
 238
 239     Ref: https://docs.pytest.org/en/latest/fixture.html#using-fixtures-from-classes-modules-or-projects
 240     """
 241     with _cleandir(remove_tempdir) as newdir:
 242         yield newdir
 243
 244
 245 @pytest.fixture(scope='session')
 246 def gmxcli():
 247     from .commandline import cli_executable
 248     command = cli_executable()
 249
 250     if command is None:
 251         message = "Tests need 'gmx' command line tool, but could not find it on the path."
 252         raise RuntimeError(message)
 253     try:
 254         assert os.access(command, os.X_OK)
 255     except Exception as E:
 256         raise RuntimeError('"{}" is not an executable gmx wrapper program'.format(command)) from E
 257     yield command