diff --git a/tests/test_flux.py b/tests/test_flux.py index 837893e6..efa3be47 100644 --- a/tests/test_flux.py +++ b/tests/test_flux.py @@ -39,7 +39,7 @@ class TestFlux(unittest.TestCase): def setUp(self): self.executor = FluxExecutor() - def test_flux_executor(self): + def test_flux_executor_serial(self): with PyFluxExecutor(max_workers=2, executor=self.executor) as exe: fs_1 = exe.submit(calc, 1) fs_2 = exe.submit(calc, 2) @@ -48,6 +48,12 @@ def test_flux_executor(self): self.assertTrue(fs_1.done()) self.assertTrue(fs_2.done()) + def test_flux_executor_parallel(self): + with PyFluxExecutor(max_workers=1, cores_per_worker=2, executor=self.executor) as exe: + fs_1 = exe.submit(mpi_funct, 1) + self.assertEqual(fs_1.result(), [(1, 2, 0), (1, 2, 1)]) + self.assertTrue(fs_1.done()) + def test_single_task(self): with SingleTaskExecutor(cores=2, executor=self.executor) as p: output = p.map(mpi_funct, [1, 2, 3]) diff --git a/tests/test_meta.py b/tests/test_meta.py index 869f9de3..81263e9d 100644 --- a/tests/test_meta.py +++ b/tests/test_meta.py @@ -15,6 +15,13 @@ def calc(i): return i +def mpi_funct(i): + from mpi4py import MPI + size = MPI.COMM_WORLD.Get_size() + rank = MPI.COMM_WORLD.Get_rank() + return i, size, rank + + class TestFutureCreation(unittest.TestCase): def test_get_future_done(self): f = _get_future_done() @@ -57,7 +64,7 @@ def test_executor_broker(self): class TestMetaExecutor(unittest.TestCase): - def test_meta_executor(self): + def test_meta_executor_serial(self): with HPCExecutor(max_workers=2) as exe: fs_1 = exe.submit(calc, 1) fs_2 = exe.submit(calc, 2) @@ -65,3 +72,9 @@ def test_meta_executor(self): self.assertEqual(fs_2.result(), 2) self.assertTrue(fs_1.done()) self.assertTrue(fs_2.done()) + + def test_meta_executor_parallel(self): + with HPCExecutor(max_workers=1, cores_per_worker=2) as exe: + fs_1 = exe.submit(mpi_funct, 1) + self.assertEqual(fs_1.result(), [(1, 2, 0), (1, 2, 1)]) + self.assertTrue(fs_1.done())