Make test_persistent_uniform_gen_decides_stop.py more robust

Libensemble · May 1, 2023 · d0a9a99 · d0a9a99
1 parent 7857645
commit d0a9a99
Showing 1 changed file with 6 additions and 9 deletions.
diff --git a/libensemble/tests/functionality_tests/test_persistent_uniform_gen_decides_stop.py b/libensemble/tests/functionality_tests/test_persistent_uniform_gen_decides_stop.py
@@ -13,7 +13,7 @@
 
 # Do not change these lines - they are parsed by run-tests.sh
 # TESTSUITE_COMMS: mpi local tcp
-# TESTSUITE_NPROCS: 5
+# TESTSUITE_NPROCS: 5, 7
 # TESTSUITE_OS_SKIP: WIN
 
 import sys
@@ -32,11 +32,9 @@
 if __name__ == "__main__":
     nworkers, is_manager, libE_specs, _ = parse_args()
 
-    if nworkers < 2:
-        sys.exit("Cannot run with a persistent worker if only one worker -- aborting...")
-
     n = 2
     ngens = 2
+    init_batch_size = nworkers - ngens
 
     if ngens >= nworkers:
         sys.exit("You number of generators must be less than the number of workers -- aborting...")
@@ -53,7 +51,7 @@
         "persis_in": ["f", "x", "sim_id"],
         "out": [("x", float, (n,))],
         "user": {
-            "initial_batch_size": nworkers,
+            "initial_batch_size": init_batch_size,
             "shutdown_limit": 10,  # Iterations on a gen before it triggers a shutdown.
             "lb": np.array([-3, -2]),
             "ub": np.array([3, 2]),
@@ -77,14 +75,13 @@
 
     if is_manager:
         [ended_times, counts] = np.unique(H["gen_ended_time"], return_counts=True)
-        print("Ended times:", ended_times)
         print("Num. points in each gen iteration:", counts)
         assert (
-            counts[0] == nworkers
+            counts[0] == init_batch_size
         ), "The first gen_ended_time should be common among initial_batch_size number of points"
         assert (
-            counts[1] == nworkers
-        ), "The second gen_ended_time should be common among initial_batch_size number of points"
+            sum(counts == init_batch_size) >= ngens
+        ), "The initial batch of each gen should be common among initial_batch_size number of points"
         assert (
             len(np.unique(counts)) > 1
         ), "All gen_ended_times are the same; they should be different for the async case"