Parameterise test-fci-mpi

tests/integrated/test-fci-mpi/test_fci_mpi.py

@@ -2,71 +2,55 @@
 #
 # Python script to run and analyse MPI test
 
-from boututils.run_wrapper import launch_safe
-from boutdata.collect import collect
-import itertools
-
 import numpy.testing as npt
+import pytest
+from boutdata.collect import collect, create_cache
+from boututils.run_wrapper import launch_safe
 
-# Resolution in x and y
-NLIST = [1, 2, 4]
 MAXCORES = 8
 NSLICES = [1]
+IMPLEMENTATIONS = ["hermitespline", "monotonichermitespline"]
+COLLECT_KW = {"info": False, "xguards": False, "yguards": False, "path": "data"}
 
 
-def test_fci_mpi():
-
-    COLLECT_KW = dict(info=False, xguards=False, yguards=False, path="data")
-
-    def run_case(nxpe: int, nype: int, mthread: int):
-
-        cmd = f"./fci_mpi NXPE={nxpe} NYPE={nype} mesh:paralleltransform:xzinterpolation:type={implementation}"
-        print(f"Running command: {cmd}")
+def run_case(nxpe: int, nype: int, mthread: int, nslice: int, implementation: str):
+    cmd = f"./fci_mpi NXPE={nxpe} NYPE={nype} mesh:paralleltransform:xzinterpolation:type={implementation}"
+    print(f"Running command: {cmd}")
 
-        _, out = launch_safe(cmd, nproc=nxpe * nype, mthread=mthread, pipe=True)
+    _, out = launch_safe(cmd, nproc=nxpe * nype, mthread=mthread, pipe=True)
 
-        # Save output to log file
-        with open(f"run.log.{nxpe}.{nype}.{nslice}.log", "w") as f:
-            f.write(out)
+    # Save output to log file
+    with open(f"run.log.{nxpe}.{nype}.{nslice}.log", "w") as f:
+        f.write(out)
 
-    def test_case(nxpe: int, nype: int, mthread: int, ref: dict) -> bool:
-        run_case(nxpe, nype, mthread)
 
-        failures = ""
+@pytest.mark.parametrize("implementation", IMPLEMENTATIONS)
+@pytest.mark.parametrize("nslice", NSLICES)
+def test_fci_mpi(implementation, nslice, subtests):
+    run_case(1, 1, MAXCORES, nslice, implementation)
+    cache = create_cache(path="data", prefix="BOUT.dmp")
 
-        for name, val in ref.items():
-            try:
-                npt.assert_allclose(val, collect(name, **COLLECT_KW))
-            except AssertionError as e:
-                failures += f"\ncase {nxpe=} {nype=} {name=}\n{e}"
-
-        return failures
-
-    failures = ""
+    ref = {}
+    for i in range(4):
+        for yp in range(1, nslice + 1):
+            for y in [-yp, yp]:
+                name = f"output_{i}_{y:+d}"
+                ref[name] = collect(name, datafile_cache=cache, **COLLECT_KW)
 
-    for implementation in ["hermitespline", "monotonichermitespline"]:
-        for nslice in NSLICES:
-            # reference data!
-            run_case(1, 1, MAXCORES)
+    # Explicitly delete cache to close netCDF files
+    del cache
 
-            ref = {}
-            for i in range(4):
-                for yp in range(1, nslice + 1):
-                    for y in [-yp, yp]:
-                        name = f"output_{i}_{y:+d}"
-                        ref[name] = collect(name, **COLLECT_KW)
+    for nxpe, nype in [(2, 1), (1, 2), (2, 2)]:
+        mthread = MAXCORES // (nxpe * nype)
+        run_case(nxpe, nype, mthread, nslice, implementation)
 
-            for nxpe, nype in itertools.product(NLIST, NLIST):
-                if (nxpe, nype) == (1, 1):
-                    # reference case, done above
-                    continue
+        cache = create_cache(path="data", prefix="BOUT.dmp")
+        results = {
+            name: collect(name, datafile_cache=cache, **COLLECT_KW) for name in ref
+        }
+        # Explicitly delete cache to close netCDF files
+        del cache
 
-                if nxpe * nype > MAXCORES:
-                    continue
-
-                mthread = MAXCORES // (nxpe * nype)
-                failures_ = test_case(nxpe, nype, mthread, ref)
-                failures += failures_
-
-        success = len(failures) == 0
-        assert success, "\nSome tests failed:\n" + failures
+        for name, val in ref.items():
+            with subtests.test(msg=f"{name}, {nxpe=}, {nype=}"):
+                npt.assert_allclose(val, results[name])