1#
2# This file is part of the GROMACS molecular simulation package.
3#
4# Copyright (c) 2019,2020, 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
37Provides utilities and pytest fixtures for gmxapi and GROMACS tests.
38
39To load these facilities in a pytest environment, set a `pytest_plugins`
40variable 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
50import json
51import logging
52import os
53import shutil
54import tempfile
55import warnings
56from contextlib import contextmanager
57from enum import Enum
58from typing import Union
59
60import pytest
61
62mpi_status = 'Test requires mpi4py managing 2 MPI ranks.'
63skip_mpi = False
64try:
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.'
73except ImportError:
74    skip_mpi = True
75    mpi_status += ' mpi4py is not available.'
76
77
78def pytest_configure(config):
79    config.addinivalue_line("markers", "withmpi_only: test requires mpi4py managing 2 MPI ranks.")
80
81
82def 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
92def 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
108class 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')
116def 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@pytest.fixture(scope='session')
122def gmxconfig():
123    try:
124        from importlib.resources import open_text
125        with open_text('gmxapi', 'gmxconfig.json') as textfile:
126            config = json.load(textfile)
127    except ImportError:
128        # TODO: Remove this when we require Python 3.7
129        try:
130            # A backport of importlib.resources is available as importlib_resources
131            # with a somewhat different interface.
132            from importlib_resources import files, as_file
133
134            source = files('gmxapi').joinpath('gmxconfig.json')
135            with as_file(source) as gmxconfig:
136                with open(gmxconfig, 'r') as fp:
137                    config = json.load(fp)
138        except ImportError:
139            config = None
140    yield config
141
142@pytest.fixture(scope='session')
143def mdrun_kwargs(request, gmxconfig):
144    """pytest fixture to provide a mdrun_kwargs dictionary for the mdrun ResourceManager.
145    """
146    from gmxapi.simulation.mdrun import ResourceManager as _ResourceManager
147    if gmxconfig is None:
148        raise RuntimeError('--threads argument requires a usable gmxconfig.json')
149    arg = request.config.getoption('--threads')
150    if arg is None:
151        return {}
152    mpi_type = gmxconfig['gmx_mpi_type']
153    if mpi_type is not None and mpi_type == "tmpi":
154        kwargs = {'threads': int(arg)}
155    else:
156        kwargs = {}
157    # TODO: (#3718) Normalize the handling of run-time arguments.
158    _ResourceManager.mdrun_kwargs = dict(**kwargs)
159    return kwargs
160
161
162@contextmanager
163def scoped_chdir(dir):
164    oldpath = os.getcwd()
165    os.chdir(dir)
166    try:
167        yield dir
168        # If the `with` block using scoped_chdir produces an exception, it will
169        # be raised at this point in this function. We want the exception to
170        # propagate out of the `with` block, but first we want to restore the
171        # original working directory, so we skip `except` but provide a `finally`.
172    finally:
173        os.chdir(oldpath)
174
175
176@contextmanager
177def _cleandir(remove_tempdir: Union[str, RmOption]):
178    """Context manager for a clean temporary working directory.
179
180    Arguments:
181        remove_tempdir (RmOption): whether to remove temporary directory "always",
182                                   "never", or on "success"
183
184    Raises:
185        ValueError: if remove_tempdir value is not valid.
186
187    The context manager will issue a warning for each temporary directory that
188    is not removed.
189    """
190    if not isinstance(remove_tempdir, RmOption):
191        remove_tempdir = RmOption(remove_tempdir)
192
193    newpath = tempfile.mkdtemp()
194
195    def remove():
196        shutil.rmtree(newpath)
197
198    def warn():
199        warnings.warn('Temporary directory not removed: {}'.format(newpath))
200
201    # Initialize callback function reference
202    if remove_tempdir == RmOption.always:
203        callback = remove
204    else:
205        callback = warn
206
207    try:
208        with scoped_chdir(newpath):
209            yield newpath
210        # If we get to this line, the `with` block using _cleandir did not throw.
211        # Clean up the temporary directory unless the user specified `--rm never`.
212        # I.e. If the user specified `--rm success`, then we need to toggle from `warn` to `remove`.
213        if remove_tempdir != RmOption.never:
214            callback = remove
215    finally:
216        callback()
217
218
219@pytest.fixture
220def cleandir(remove_tempdir: RmOption):
221    """Provide a clean temporary working directory for a test.
222
223    Example usage:
224
225        import os
226        import pytest
227
228        @pytest.mark.usefixtures("cleandir")
229        def test_cwd_starts_empty():
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(cleandir):
235            assert os.listdir(os.getcwd()) == []
236            assert os.path.abspath(os.getcwd()) == os.path.abspath(cleandir)
237            with open("myfile", "w") as f:
238                f.write("hello")
239
240        @pytest.mark.usefixtures("cleandir")
241        class TestDirectoryInit(object):
242            def test_cwd_starts_empty(self):
243                assert os.listdir(os.getcwd()) == []
244                with open("myfile", "w") as f:
245                    f.write("hello")
246
247            def test_cwd_also_starts_empty(self):
248                assert os.listdir(os.getcwd()) == []
249                with open("myfile", "w") as f:
250                    f.write("hello")
251
252    Ref: https://docs.pytest.org/en/latest/fixture.html#using-fixtures-from-classes-modules-or-projects
253    """
254    with _cleandir(remove_tempdir) as newdir:
255        yield newdir
256
257
258class GmxBin:
259    """Represent the detected GROMACS installation."""
260    def __init__(self, gmxconfig):
261        # Try to use package resources to locate the "gmx" binary wrapper.
262        if gmxconfig is not None:
263            gmxbindir = gmxconfig.get('gmx_bindir', None)
264            command = gmxconfig.get('gmx_executable', None)
265        else:
266            gmxbindir = None
267            command = None
268
269        # TODO: Remove fall-back when we can rely on gmxconfig.json via importlib.resources in Py 3.7+.
270        allowed_command_names = ['gmx', 'gmx_mpi']
271        for command_name in allowed_command_names:
272            if command is not None:
273                break
274            command = shutil.which(command_name)
275            if command is None:
276                gmxbindir = os.getenv('GMXBIN')
277                if gmxbindir is None:
278                    gromacsdir = os.getenv('GROMACS_DIR')
279                    if gromacsdir is not None and gromacsdir != '':
280                        gmxbindir = os.path.join(gromacsdir, 'bin')
281                if gmxbindir is None:
282                    gmxapidir = os.getenv('gmxapi_DIR')
283                    if gmxapidir is not None and gmxapidir != '':
284                        gmxbindir = os.path.join(gmxapidir, 'bin')
285                if gmxbindir is not None:
286                    gmxbindir = os.path.abspath(gmxbindir)
287                    command = shutil.which(command_name, path=gmxbindir)
288
289        self._command = command
290        self._bindir = gmxbindir
291
292    def command(self):
293        return self._command
294
295    def bindir(self):
296        return self._bindir
297
298
299@pytest.fixture(scope='session')
300def gmxcli(gmxconfig):
301    command = GmxBin(gmxconfig).command()
302    if command is None:
303        message = "Tests need 'gmx' command line tool, but could not find it on the path."
304        raise RuntimeError(message)
305    try:
306        assert os.access(command, os.X_OK)
307    except Exception as E:
308        raise RuntimeError('"{}" is not an executable gmx wrapper program'.format(command)) from E
309    yield command
310