From d94eb4373816225582ba922b8d9586b9cd19da59 Mon Sep 17 00:00:00 2001
From: Matthias Diener <mdiener@illinois.edu>
Date: Sun, 4 Dec 2022 15:52:46 -0600
Subject: [PATCH] add pytest_raises_on_rank (#162)
* add pytest_raises_on_rank
* better type annotation
* add copyright to mpi.py
* add test_mpi.py
* fix doc
* Clarify test
---
doc/conf.py | 3 +++
doc/index.rst | 1 +
doc/mpi.rst | 1 +
pytools/mpi.py | 58 ++++++++++++++++++++++++++++++++++++++++++++++++
test/test_mpi.py | 46 ++++++++++++++++++++++++++++++++++++++
5 files changed, 109 insertions(+)
create mode 100644 doc/mpi.rst
create mode 100644 test/test_mpi.py
diff --git a/doc/conf.py b/doc/conf.py
index 70a55f0..1338441 100644
--- a/doc/conf.py
+++ b/doc/conf.py
@@ -29,8 +29,11 @@ intersphinx_mapping = {
"https://numpy.org/doc/stable": None,
"https://documen.tician.de/pymbolic/": None,
"https://documen.tician.de/loopy/": None,
+ "https://docs.pytest.org/en/stable/": None,
}
nitpick_ignore_regex = [
["py:class", r"typing_extensions\.(.+)"],
]
+
+nitpicky = True
diff --git a/doc/index.rst b/doc/index.rst
index 25f26c6..1a1e9c4 100644
--- a/doc/index.rst
+++ b/doc/index.rst
@@ -11,6 +11,7 @@ Welcome to pytools's documentation!
graph
tag
codegen
+ mpi
misc
🚀 Github <https://github.com/inducer/pytools>
💾 Download Releases <https://pypi.python.org/pypi/pytools>
diff --git a/doc/mpi.rst b/doc/mpi.rst
new file mode 100644
index 0000000..e7bcae5
--- /dev/null
+++ b/doc/mpi.rst
@@ -0,0 +1 @@
+.. automodule:: pytools.mpi
diff --git a/pytools/mpi.py b/pytools/mpi.py
index e165019..9961d76 100644
--- a/pytools/mpi.py
+++ b/pytools/mpi.py
@@ -1,3 +1,41 @@
+__copyright__ = """
+Copyright (C) 2009-2019 Andreas Kloeckner
+Copyright (C) 2022 University of Illinois Board of Trustees
+"""
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+__doc__ = """
+MPI helper functionality
+========================
+
+.. autofunction:: check_for_mpi_relaunch
+.. autofunction:: run_with_mpi_ranks
+.. autofunction:: pytest_raises_on_rank
+"""
+
+from contextlib import contextmanager, AbstractContextManager
+from typing import Generator, Tuple, Union, Type
+
+
def check_for_mpi_relaunch(argv):
if argv[1] != "--mpi-relaunch":
return
@@ -26,3 +64,23 @@ def run_with_mpi_ranks(py_script, ranks, callable_, args=(), kwargs=None):
check_call(["mpirun", "-np", str(ranks),
sys.executable, py_script, "--mpi-relaunch", callable_and_args],
env=newenv)
+
+
+@contextmanager
+def pytest_raises_on_rank(my_rank: int, fail_rank: int,
+ expected_exception: Union[Type[BaseException],
+ Tuple[Type[BaseException], ...]]) \
+ -> Generator[AbstractContextManager, None, None]:
+ """
+ Like :func:`pytest.raises`, but only expect an exception on rank *fail_rank*.
+ """
+ import pytest
+ from contextlib import nullcontext
+
+ if my_rank == fail_rank:
+ cm: AbstractContextManager = pytest.raises(expected_exception)
+ else:
+ cm = nullcontext()
+
+ with cm as exc:
+ yield exc
diff --git a/test/test_mpi.py b/test/test_mpi.py
new file mode 100644
index 0000000..4bbd347
--- /dev/null
+++ b/test/test_mpi.py
@@ -0,0 +1,46 @@
+__copyright__ = "Copyright (C) 2022 University of Illinois Board of Trustees"
+
+__license__ = """
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
+"""
+
+import pytest
+
+
+def test_pytest_raises_on_rank():
+ from pytools.mpi import pytest_raises_on_rank
+
+ def fail(my_rank: int, fail_rank: int) -> None:
+ if my_rank == fail_rank:
+ raise ValueError("test failure")
+
+ with pytest.raises(ValueError):
+ fail(0, 0)
+
+ fail(0, 1)
+
+ with pytest_raises_on_rank(0, 0, ValueError):
+ # Generates an exception, and pytest_raises_on_rank
+ # expects one.
+ fail(0, 0)
+
+ with pytest_raises_on_rank(0, 1, ValueError):
+ # Generates no exception, and pytest_raises_on_rank
+ # does not expect one.
+ fail(0, 1)
--
GitLab