From 1c43320cd902f5caa72c77370840f56a32933f20 Mon Sep 17 00:00:00 2001
From: Kevin Paul <kpaul@nvidia.com>
Date: Thu, 12 Oct 2023 15:51:46 +0200
Subject: [PATCH] Rename test / add no_exit test for execute

---
 dask_mpi/tests/execute_no_exit.py | 30 ++++++++++++++++++++++++++++++
 dask_mpi/tests/test_execute.py    | 15 +++++++++++++--
 2 files changed, 43 insertions(+), 2 deletions(-)
 create mode 100644 dask_mpi/tests/execute_no_exit.py

diff --git a/dask_mpi/tests/execute_no_exit.py b/dask_mpi/tests/execute_no_exit.py
new file mode 100644
index 0000000..548f227
--- /dev/null
+++ b/dask_mpi/tests/execute_no_exit.py
@@ -0,0 +1,30 @@
+from time import time, sleep
+
+from distributed import Client
+from mpi4py.MPI import COMM_WORLD as world
+
+from dask_mpi import execute
+
+# Split our MPI world into two pieces, one consisting just of
+# the old rank 3 process and the other with everything else
+new_comm_assignment = 1 if world.rank == 3 else 0
+comm = world.Split(new_comm_assignment)
+
+if world.rank != 3:
+    def client_code():
+        with Client() as c:
+            start = time()
+            while len(c.scheduler_info()["workers"]) != 1:
+                assert time() < start + 10
+                sleep(0.2)
+
+            c.submit(lambda x: x + 1, 10).result() == 11
+            c.submit(lambda x: x + 1, 20).result() == 21
+
+    execute(client_code, comm=comm)
+
+# check that our original comm is intact
+world.Barrier()
+x = 100 if world.rank == 0 else 200
+x = world.bcast(x)
+assert x == 100
diff --git a/dask_mpi/tests/test_execute.py b/dask_mpi/tests/test_execute.py
index ab2ff08..318d10e 100644
--- a/dask_mpi/tests/test_execute.py
+++ b/dask_mpi/tests/test_execute.py
@@ -22,9 +22,9 @@
         (1, ["-c", "0", "-s", "0", "-x", "True"], 1),
     ],
 )
-def test_execute(mpisize, execute_args, retcode, mpirun):
+def test_basic(mpisize, execute_args, retcode, mpirun):
     script_file = os.path.join(
-        os.path.dirname(os.path.realpath(__file__)), "execute_script.py"
+        os.path.dirname(os.path.realpath(__file__)), "execute_basic.py"
     )
 
     execute_args += ["-m", str(mpisize)]
@@ -34,3 +34,14 @@ def test_execute(mpisize, execute_args, retcode, mpirun):
 
     p.communicate()
     assert p.returncode == retcode
+
+
+def test_no_exit(mpirun):
+    script_file = os.path.join(
+        os.path.dirname(os.path.realpath(__file__)), "execute_no_exit.py"
+    )
+
+    p = subprocess.Popen(mpirun + ["-np", "4", sys.executable, script_file])
+
+    p.communicate()
+    assert p.returncode == 0