map(): Evaluate generator within with-statement (#812)

Author: jan-janssen

tests/test_fluxjobexecutor.py

@@ -104,9 +104,9 @@ def test_single_task(self):
             block_allocation=True,
             pmi_mode=pmi,
         ) as p:
-            output = p.map(mpi_funct, [1, 2, 3])
+            output = list(p.map(mpi_funct, [1, 2, 3]))
         self.assertEqual(
-            list(output),
+            output,
             [[(1, 2, 0), (1, 2, 1)], [(2, 2, 0), (2, 2, 1)], [(3, 2, 0), (3, 2, 1)]],
         )
 
@@ -122,9 +122,9 @@ def test_output_files_cwd(self):
             block_allocation=True,
             flux_log_files=True,
         ) as p:
-            output = p.map(calc, [1, 2, 3])
+            output = list(p.map(calc, [1, 2, 3]))
         self.assertEqual(
-            list(output),
+            output,
             [1, 2, 3],
         )
         self.assertTrue(os.path.exists(file_stdout))
@@ -142,9 +142,9 @@ def test_output_files_abs(self):
             block_allocation=True,
             flux_log_files=True,
         ) as p:
-            output = p.map(calc, [1, 2, 3])
+            output = list(p.map(calc, [1, 2, 3]))
         self.assertEqual(
-            list(output),
+            output,
             [1, 2, 3],
         )
         self.assertTrue(os.path.exists(file_stdout))

tests/test_fluxpythonspawner.py

@@ -99,9 +99,9 @@ def test_single_task(self):
             },
             spawner=FluxPythonSpawner,
         ) as p:
-            output = p.map(mpi_funct, [1, 2, 3])
+            output = list(p.map(mpi_funct, [1, 2, 3]))
         self.assertEqual(
-            list(output),
+            output,
             [[(1, 2, 0), (1, 2, 1)], [(2, 2, 0), (2, 2, 1)], [(3, 2, 0), (3, 2, 1)]],
         )
 

tests/test_mpiexecspawner.py

@@ -323,8 +323,8 @@ def test_pool_serial_map(self):
             executor_kwargs={"cores": 1},
             spawner=MpiExecSpawner,
         ) as p:
-            output = p.map(calc_array, [1, 2, 3])
-        self.assertEqual(list(output), [np.array(1), np.array(4), np.array(9)])
+            output = list(p.map(calc_array, [1, 2, 3]))
+        self.assertEqual(output, [np.array(1), np.array(4), np.array(9)])
 
     def test_executor_exception(self):
         with self.assertRaises(RuntimeError):
@@ -430,9 +430,9 @@ def test_pool_multi_core_map(self):
             executor_kwargs={"cores": 2},
             spawner=MpiExecSpawner,
         ) as p:
-            output = p.map(mpi_funct, [1, 2, 3])
+            output = list(p.map(mpi_funct, [1, 2, 3]))
         self.assertEqual(
-            list(output),
+            output,
             [[(1, 2, 0), (1, 2, 1)], [(2, 2, 0), (2, 2, 1)], [(3, 2, 0), (3, 2, 1)]],
         )
 

tests/test_singlenodeexecutor_mpi.py

@@ -119,9 +119,9 @@ def test_output_files_cwd(self):
             resource_dict={"cores": 1, "cwd": dirname},
             block_allocation=True,
         ) as p:
-            output = p.map(calc, [1, 2, 3])
+            output = list(p.map(calc, [1, 2, 3]))
         self.assertEqual(
-            list(output),
+            output,
             [1, 2, 3],
         )