Create mixins for different test functionality

reframe_tests/tests/base_tests.py

@@ -7,11 +7,9 @@ Usage (run all tests):
 
 import os
 from pathlib import Path
-from shutil import copyfile
+from typing import Literal, Optional, Union
-from typing import Literal
 
 import reframe.utility.sanity as sn
-import reframe.utility.typecheck as typ
 from numpy import prod
 from reframe import RunOnlyRegressionTest, simple_test
 from reframe.core.builtins import (
@@ -22,11 +20,10 @@ from reframe.core.builtins import (
     run_after,
     run_before,
     sanity_function,
-    variable,
 )
-from reframe.utility import osext, udeps
+from reframe.utility import udeps
 
-from gprMax.receivers import Rx
+from reframe_tests.tests.regression_checks import RegressionCheck
 from reframe_tests.utilities.deferrable import path_join
 
 TESTS_ROOT_DIR = Path(__file__).parent
@@ -106,14 +103,30 @@ class GprMaxRegressionTest(RunOnlyRegressionTest):
     exclusive_access = True
 
     model = parameter()
-    is_antenna_model = variable(bool, value=False)
-    has_receiver_output = variable(bool, value=True)
-    snapshots = variable(typ.List[str], value=[])
     sourcesdir = required
-    extra_executable_opts = variable(typ.List[str], value=[])
     executable = "time -p python -m gprMax --log-level 10 --hide-progress-bars"
 
-    rx_outputs = variable(typ.List[str], value=Rx.defaultoutputs)
+    regression_checks: list[RegressionCheck] = []
+
+    test_dependency: Optional[type["GprMaxRegressionTest"]] = None
+
+    def get_test_dependency(self) -> Optional["GprMaxRegressionTest"]:
+        """Get test variant with the same model and number of models"""
+        if self.test_dependency is None:
+            return None
+        else:
+            variant = self.test_dependency.variant_name(self.test_dependency.param_variant)
+            return self.getdep(variant)
+
+    def build_reference_filepath(self, name: Union[str, os.PathLike]) -> Path:
+        target = self.get_test_dependency()
+        if target is None:
+            reference_dir = self.short_name
+        else:
+            reference_dir = target.short_name
+
+        reference_file = Path("regression_checks", reference_dir, name).with_suffix(".h5")
+        return reference_file.absolute()
 
     @run_after("init")
     def setup_env_vars(self):
@@ -133,6 +146,9 @@ class GprMaxRegressionTest(RunOnlyRegressionTest):
     def inject_dependencies(self):
         """Test depends on the Python virtual environment building correctly"""
         self.depends_on("CreatePyenvTest", udeps.by_env)
+        if self.test_dependency is not None:
+            variant = self.test_dependency.variant_name(self.test_dependency.param_variant)
+            self.depends_on(variant, udeps.by_env)
 
     @require_deps
     def get_pyenv_path(self, CreatePyenvTest):
@@ -140,35 +156,19 @@ class GprMaxRegressionTest(RunOnlyRegressionTest):
         path_to_pyenv = os.path.join(CreatePyenvTest(part="login").stagedir, PATH_TO_PYENV)
         self.prerun_cmds.append(f"source {path_to_pyenv}")
 
-    def build_reference_filepath(self, suffix: str = "") -> str:
+    @run_after("init")
-        filename = f"{self.short_name}_{suffix}" if len(suffix) > 0 else self.short_name
+    def configure_test_run(self):
-        reference_file = Path("regression_checks", filename).with_suffix(".h5")
-        return os.path.abspath(reference_file)
-
-    def build_snapshot_filepath(self, snapshot: str) -> str:
-        return os.path.join(f"{self.model}_snaps", snapshot)
-
-    @run_after("setup")
-    def setup_reference_files(self):
-        """Build reference file paths"""
-        self.reference_file = self.build_reference_filepath()
-        self.snapshot_reference_files = []
-        for snapshot in self.snapshots:
-            self.snapshot_reference_files.append(self.build_reference_filepath(snapshot))
-
-    @run_after("setup", always_last=True)
-    def configure_test_run(self, input_file_ext: str = ".in"):
         """Configure gprMax commandline arguments and plot outputs
 
         Set the input and output files and add postrun commands to plot
         the outputs.
         """
-        self.input_file = f"{self.model}{input_file_ext}"
+        self.input_file = f"{self.model}.in"
         self.output_file = f"{self.model}.h5"
         self.executable_opts = [self.input_file, "-o", self.output_file]
-        self.executable_opts += self.extra_executable_opts
+        self.keep_files = [self.input_file, self.output_file]
-        self.keep_files = [self.input_file, *self.snapshots]
 
+        """
         if self.has_receiver_output:
             self.postrun_cmds = [
                 f"python -m reframe_tests.utilities.plotting {self.output_file} {self.reference_file} -m {self.model}"
@@ -186,6 +186,7 @@ class GprMaxRegressionTest(RunOnlyRegressionTest):
                 antenna_t1_params,
                 antenna_ant_params,
             ]
+        """
 
     @run_before("run")
     def combine_task_outputs(self):
@@ -233,99 +234,30 @@ class GprMaxRegressionTest(RunOnlyRegressionTest):
             r"=== Simulation completed in ", self.stdout, "Simulation did not complete"
         )
 
-    def run_regression_check(
-        self, output_file: str, reference_file: str, error_msg: str
-    ) -> Literal[True]:
-        """Compare two provided .h5 files using h5diff
-
-        Args:
-            output_file: Filepath of .h5 file output by the test.
-            reference_file: Filepath of reference .h5 file containing
-                the expected output.
-        """
-        if self.current_system.name == "archer2":
-            h5diff = "/opt/cray/pe/hdf5/default/bin/h5diff"
-        else:
-            h5diff = "h5diff"
-
-        h5diff_output = osext.run_command([h5diff, os.path.abspath(output_file), reference_file])
-
-        return sn.assert_false(
-            h5diff_output.stdout,
-            (
-                f"{error_msg}\n"
-                f"For more details run: 'h5diff {os.path.abspath(output_file)} {reference_file}'\n"
-                f"To re-create regression file, delete '{reference_file}' and rerun the test."
-            ),
-        )
-
-    def test_output_regression(self) -> Literal[True]:
-        """Compare the test output with the reference file.
-
-        If the test contains any receivers, a regression check is run,
-        otherwise it checks the test did not generate an output file.
-        """
-        if self.has_receiver_output:
-            return self.run_regression_check(
-                self.output_file, self.reference_file, "Failed output regresssion check"
-            )
-        else:
-            return sn.assert_false(
-                sn.path_exists(self.output_file),
-                f"Unexpected output file found: '{self.output_file}'",
-            )
-
-    def test_snapshot_regression(self) -> Literal[True]:
-        """Compare the snapshot outputs with reference files.
-
-        Generates a regression check for each snapshot. Each regression
-        check is a deffered expression, so they all need to be returned
-        so that they are each evaluated.
-        """
-        regression_checks = []
-        for index, snapshot in enumerate(self.snapshots):
-            snapshot_path = self.build_snapshot_filepath(snapshot)
-            reference_file = self.snapshot_reference_files[index]
-            regression_checks.append(
-                self.run_regression_check(
-                    snapshot_path, reference_file, f"Failed snapshot regresssion check '{snapshot}'"
-                )
-            )
-
-        # sn.assert_true is not strictly necessary
-        return sn.assert_true(sn.all(regression_checks))
-
     @sanity_function
-    def regression_check(self) -> Literal[True]:
+    def regression_check(self) -> bool:
         """Perform regression check for the test output and snapshots
 
         If not all the reference files exist, then create all the
         missing reference files from the test output and fail the test.
         """
-        if (not self.has_receiver_output or sn.path_exists(self.reference_file)) and sn.all(
+        error_messages = []
-            [sn.path_exists(path) for path in self.snapshot_reference_files]
+        for check in self.regression_checks:
-        ):
+            if not check.reference_file_exists():
-            return (
+                if check.create_reference_file():
-                self.test_simulation_complete()
+                    error_messages.append(
-                and self.test_output_regression()
+                        f"Reference file does not exist. Creating... '{check.reference_file}'"
-                and self.test_snapshot_regression()
                     )
                 else:
-            error_messages = []
-            if self.has_receiver_output and not sn.path_exists(self.reference_file):
-                copyfile(self.output_file, self.reference_file)
                     error_messages.append(
-                    f"Output reference file does not exist. Creating... '{self.reference_file}'"
+                        f"ERROR: Unable to create reference file: '{check.reference_file}'"
-                )
-            for index, snapshot in enumerate(self.snapshots):
-                reference_file = self.snapshot_reference_files[index]
-                if not sn.path_exists(reference_file):
-                    copyfile(self.build_snapshot_filepath(snapshot), reference_file)
-                    error_messages.append(
-                        f"Snapshot '{snapshot}' reference file does not exist. Creating... '{reference_file}'"
                     )
 
-            return sn.assert_true(False, "\n".join(error_messages))
+        return (
+            self.test_simulation_complete()
+            and sn.assert_true(len(error_messages) < 1, "\n".join(error_messages))
+            and sn.all(sn.map(lambda check: check.run(), self.regression_checks))
+        )
 
     @performance_function("s", perf_key="run_time")
     def extract_run_time(self):

reframe_tests/tests/mixins.py

@@ -0,0 +1,114 @@
+from pathlib import Path
+
+import reframe.utility.typecheck as typ
+from numpy import prod
+from reframe import RegressionMixin
+from reframe.core.builtins import parameter, required, run_after, variable
+from typing_extensions import TYPE_CHECKING
+
+from reframe_tests.tests.base_tests import GprMaxRegressionTest
+from reframe_tests.tests.regression_checks import (
+    ReceiverRegressionCheck,
+    RegressionCheck,
+    SnapshotRegressionCheck,
+)
+
+if TYPE_CHECKING:
+    GprMaxMixin = GprMaxRegressionTest
+else:
+    GprMaxMixin = RegressionMixin
+
+
+class ReceiverMixin(GprMaxMixin):
+    number_of_receivers = variable(int, value=-1)
+
+    @run_after("setup")
+    def add_receiver_regression_checks(self):
+        reference_file = self.build_reference_filepath(self.output_file)
+
+        if self.number_of_receivers > 0:
+            for i in range(self.number_of_receivers):
+                regression_check = ReceiverRegressionCheck(
+                    self.output_file, reference_file, f"r{i}"
+                )
+                self.regression_checks.append(regression_check)
+        else:
+            regression_check = RegressionCheck(self.output_file, reference_file)
+            self.regression_checks.append(regression_check)
+
+
+class SnapshotMixin(GprMaxMixin):
+    snapshots = variable(typ.List[str], value=[])
+
+    def build_snapshot_filepath(self, snapshot: str) -> Path:
+        return Path(f"{self.model}_snaps", snapshot).with_suffix(".h5")
+
+    @run_after("setup")
+    def add_snapshot_regression_checks(self):
+        has_specified_snapshots = len(self.snapshots) > 0
+        valid_test_dependency = self.test_dependency is not None and issubclass(
+            self.test_dependency, SnapshotMixin
+        )
+
+        self.skip_if(
+            not valid_test_dependency and not has_specified_snapshots,
+            f"Must provide either a list of snapshots, or a test dependency that inherits from SnapshotMixin.",
+        )
+        self.skip_if(
+            valid_test_dependency and has_specified_snapshots,
+            f"Cannot provide both a list of snapshots, and a test dependency that inherits from SnapshotMixin.",
+        )
+
+        if valid_test_dependency:
+            target = self.get_test_dependency()
+            assert isinstance(target, SnapshotMixin)
+            self.snapshots = target.snapshots
+
+        for snapshot in self.snapshots:
+            snapshot_file = self.build_snapshot_filepath(snapshot)
+            reference_file = self.build_reference_filepath(snapshot)
+            regression_check = SnapshotRegressionCheck(snapshot_file, reference_file)
+            self.regression_checks.append(regression_check)
+
+
+class PythonApiMixin(GprMaxMixin):
+    executable = "time -p python"
+
+    @run_after("setup")
+    def set_python_input_file(self):
+        """Input files for API tests will be python files"""
+        self.input_file = self.input_file.with_suffix(".py")
+
+
+class MpiMixin(GprMaxMixin):
+    mpi_layout = parameter()
+
+    @run_after("setup")
+    def configure_mpi_tasks(self):
+        """Add MPI specific commandline arguments"""
+        self.num_tasks = int(prod(self.mpi_layout))
+        self.executable_opts += ["--mpi", *map(str, self.mpi_layout)]
+
+
+class BScanMixin(GprMaxMixin):
+    num_models = parameter()
+
+    @run_after("setup")
+    def setup_bscan_test(self):
+        """Add B-Scan specific commandline arguments and postrun cmds"""
+        self.executable_opts += ["-n", str(self.num_models)]
+
+
+class TaskfarmMixin(GprMaxMixin):
+    extra_executable_opts = ["-taskfarm"]
+
+    num_tasks = required
+
+    @run_after("setup")
+    def add_taskfarm_flag(self):
+        """Add taskfarm specific commandline arguments"""
+        self.executable_opts += ["-taskfarm"]
+
+
+class AntennaModelMixin(GprMaxMixin):
+    pass

reframe_tests/tests/regression_checks.py

@@ -0,0 +1,85 @@
+from os import PathLike
+from pathlib import Path
+from shutil import copyfile
+from typing import Literal, Optional, Union
+
+import reframe.utility.sanity as sn
+from reframe.core.runtime import runtime
+from reframe.utility import osext
+
+
+class RegressionCheck:
+    """Compare two .h5 files using h5diff"""
+
+    def __init__(
+        self, output_file: Union[str, PathLike], reference_file: Union[str, PathLike]
+    ) -> None:
+        self.output_file = Path(output_file)
+        self.reference_file = Path(reference_file)
+        self.h5diff_options: list[str] = []
+
+    @property
+    def error_msg(self) -> str:
+        return "Failed regression check"
+
+    def create_reference_file(self) -> bool:
+        if not sn.path_exists(self.reference_file):
+            self.reference_file.parent.mkdir(parents=True, exist_ok=True)
+            copyfile(self.output_file, self.reference_file)
+            return True
+        else:
+            return False
+
+    def reference_file_exists(self) -> bool:
+        return sn.path_isfile(self.reference_file)
+
+    def run(self) -> Literal[True]:
+        if runtime().system.name == "archer2":
+            h5diff = "/opt/cray/pe/hdf5/default/bin/h5diff"
+        else:
+            h5diff = "h5diff"
+
+        h5diff_output = osext.run_command(
+            [h5diff, *self.h5diff_options, str(self.output_file), str(self.reference_file)]
+        )
+
+        return sn.assert_true(
+            sn.path_isfile(self.output_file),
+            f"Expected output file '{self.output_file}' does not exist",
+        ) and sn.assert_false(
+            h5diff_output.stdout,
+            (
+                f"{self.error_msg}\n"
+                # f"For more details run: 'h5diff {' '.join(self.h5diff_options)} {self.output_file} {self.reference_file}'\n"
+                f"For more details run: '{' '.join(h5diff_output.args)}'\n"
+                f"To re-create regression file, delete '{self.reference_file}' and rerun the test."
+            ),
+        )
+
+
+class ReceiverRegressionCheck(RegressionCheck):
+    def __init__(
+        self,
+        output_file: Union[str, PathLike],
+        reference_file: Union[str, PathLike],
+        output_receiver: Optional[str],
+        reference_receiver: Optional[str] = None,
+    ) -> None:
+        super().__init__(output_file, reference_file)
+
+        self.output_receiver = output_receiver
+        self.reference_receiver = reference_receiver
+
+        self.h5diff_options.append(f"rxs/{self.output_receiver}")
+        if self.reference_receiver is not None:
+            self.h5diff_options.append(f"rxs/{self.reference_receiver}")
+
+    @property
+    def error_msg(self) -> str:
+        return f"Receiver '{self.output_receiver}' failed regression check"
+
+
+class SnapshotRegressionCheck(RegressionCheck):
+    @property
+    def error_msg(self) -> str:
+        return f"Snapshot '{self.output_file.name}' failed regression check "

reframe_tests/tests/test_example_models.py

@@ -2,13 +2,14 @@ import reframe as rfm
 from reframe.core.builtins import parameter
 
 from reframe_tests.tests.base_tests import GprMaxBScanRegressionTest, GprMaxRegressionTest
+from reframe_tests.tests.mixins import MpiMixin, ReceiverMixin
 
 """Reframe regression tests for example models in gprMax documentation
 """
 
 
 @rfm.simple_test
-class TestAscan(GprMaxRegressionTest):
+class TestAscan(ReceiverMixin, GprMaxRegressionTest):
     tags = {
         "test",
         "serial",