diff --git a/tests/simulator/test_python_concurrent.py b/tests/simulator/test_python_concurrent.py index 89e0d9fa4..91be34ef4 100644 --- a/tests/simulator/test_python_concurrent.py +++ b/tests/simulator/test_python_concurrent.py @@ -4,54 +4,80 @@ """ import os +import time import unittest -import multiprocessing +from multiprocessing import cpu_count from tests.simulator.test_validation import SimulatorValidation +from tests.simulator.test_diagnostics import DiagnosticsTest # mpirun -n 1/2/3/4 -from tests.simulator.initialize.test_fields_init_1d import InitializationTest as InitField1d -from tests.simulator.initialize.test_particles_init_1d import InitializationTest as InitParticles1d +import tests.simulator.initialize.test_fields_init_1d as test_fields_init_1d +# from tests.simulator.initialize.test_fields_init_1d import ( +# InitializationTest as InitField1d, +# ) +# from tests.simulator.initialize.test_particles_init_1d import ( +# InitializationTest as InitParticles1d, +# ) +# from tests.simulator.advance.test_fields_advance_1d import AdvanceTest as AdvanceField1d +# from tests.simulator.advance.test_particles_advance_1d import ( +# AdvanceTest as AdvanceParticles1d, +# ) +# from tests.simulator.initialize.test_fields_init_2d import ( +# InitializationTest as InitField2d, +# ) +# from tests.simulator.initialize.test_particles_init_2d import ( +# InitializationTest as InitParticles2d, +# ) +# from tests.simulator.advance.test_fields_advance_2d import AdvanceTest as AdvanceField2d +# from tests.simulator.advance.test_particles_advance_2d import ( +# AdvanceTest as AdvanceParticles2d, +# ) -from tests.simulator.advance.test_fields_advance_1d import AdvanceTest as AdvanceField1d -from tests.simulator.advance.test_particles_advance_1d import AdvanceTest as AdvanceParticles1d +from phlop.testing.parallel_processor import load_test_cases_in, process -from tests.simulator.initialize.test_fields_init_2d import InitializationTest as InitField2d -from tests.simulator.initialize.test_particles_init_2d import InitializationTest as InitParticles2d +N_CORES = int(os.environ["N_CORES"]) if "N_CORES" in os.environ else cpu_count() +PRINT = int(os.environ["PRINT"]) if "PRINT" in os.environ else 0 +MPI_RUN = os.environ.get("MPI_RUN", None) +MPI_RUN_EXTRA = os.environ.get("MPI_RUN_EXTRA", "") -from tests.simulator.advance.test_fields_advance_2d import AdvanceTest as AdvanceField2d -from tests.simulator.advance.test_particles_advance_2d import AdvanceTest as AdvanceParticles2d +def test_cmd(clazz, test_id, mpi_run): + return f"mpirun {MPI_RUN_EXTRA} -n {mpi_run} python3 -m {clazz.__module__} {clazz.__name__}.{test_id}" -N_CORES = int(os.environ["N_CORES"]) if "N_CORES" in os.environ else multiprocessing.cpu_count() -MPI_RUN = int(os.environ["MPI_RUN"]) if "MPI_RUN" in os.environ else 1 -PRINT = int(os.environ["PRINT"]) if "PRINT" in os.environ else 0 -def test_cmd(clazz, test_id): - return f"mpirun -n {MPI_RUN} python3 -m {clazz.__module__} {clazz.__name__}.{test_id}" +test_classes_to_run = [ + SimulatorValidation, + DiagnosticsTest, + test_fields_init_1d.InitializationTest, + # InitField1d, + # InitParticles1d, + # AdvanceField1d, + # AdvanceParticles1d, + # InitField2d, + # InitParticles2d, + # AdvanceField2d, + # AdvanceParticles2d, +] -if __name__ == "__main__": - test_classes_to_run = [ - SimulatorValidation, - InitField1d, - InitParticles1d, - AdvanceField1d, - AdvanceParticles1d, - InitField2d, - InitParticles2d, - AdvanceField2d, - AdvanceParticles2d - ] - - tests = [] - loader = unittest.TestLoader() - for test_class in test_classes_to_run: - for suite in loader.loadTestsFromTestCase(test_class): - tests += [test_cmd(type(suite), suite._testMethodName)] - - from tools.python3 import run_mp - if PRINT: - for test in tests: - print(test) +def load_test_cases(classes, cores): + return load_test_cases_in(classes, cores, test_cmd_fn=test_cmd) + +def build_batches(): + batches = [] + if not MPI_RUN: + batches += [load_test_cases(test_classes_to_run, 1)] + batches += [load_test_cases(test_classes_to_run, 2)] + batches += [load_test_cases([DiagnosticsTest], 3)] + batches += [load_test_cases([DiagnosticsTest], 4)] else: - run_mp(tests, N_CORES, check=True) + batches += [load_test_cases(test_classes_to_run, int(MPI_RUN))] + return batches + + +def main(): + process(build_batches()) + + +if __name__ == "__main__": + main()