From c602fbdb559ef8aa8f065dbd7bc565c82f7e740a Mon Sep 17 00:00:00 2001
From: Matthew Smith <mjsmith6@illinois.edu>
Date: Wed, 18 Jan 2023 10:55:59 -0600
Subject: [PATCH] add test for dag with a receive as an output

---
 test/test_distributed.py | 47 ++++++++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/test/test_distributed.py b/test/test_distributed.py
index 34c1b46..81468f2 100644
--- a/test/test_distributed.py
+++ b/test/test_distributed.py
@@ -291,6 +291,53 @@ def test_dag_with_duplicated_output_arrays():
 # }}}
 
 
+# {{{ test DAG with a receive as an output
+
+def _test_dag_with_recv_as_output_inner(ctx_factory):
+    from numpy.random import default_rng
+    from mpi4py import MPI  # pylint: disable=import-error
+    comm = MPI.COMM_WORLD
+    ctx = ctx_factory()
+    queue = cl.CommandQueue(ctx)
+
+    # {{{ construct the DAG
+
+    if comm.rank == 0:
+        rng = default_rng()
+        x_np = rng.random((10, 4))
+        x = pt.make_data_wrapper(cla.to_device(queue, x_np))
+        y = 2 * x
+        send = pt.staple_distributed_send(
+            y, dest_rank=1, comm_tag=42,
+            stapled_to=pt.ones(10))
+        dag = pt.make_dict_of_named_arrays({"y": y, "send": send})
+    else:
+        y = pt.make_distributed_recv(
+            src_rank=0, comm_tag=42,
+            shape=(10, 4), dtype=np.float64)
+        dag = pt.make_dict_of_named_arrays({"y": y})
+
+    # }}}
+
+    parts = pt.find_distributed_partition(comm, dag)
+    prg_per_partition = pt.generate_code_for_partition(parts)
+    out_dict = pt.execute_distributed_partition(
+            parts, prg_per_partition, queue, comm)
+
+    if comm.rank == 0:
+        comm.bcast(x_np)
+    else:
+        x_np = comm.bcast(None)
+
+    np.testing.assert_allclose(out_dict["y"].get(), 2 * x_np)
+
+
+def test_dag_with_recv_as_output():
+    run_test_with_mpi(2, _test_dag_with_recv_as_output_inner)
+
+# }}}
+
+
 # {{{ test deterministic partitioning
 
 def _gather_random_dist_partitions(ctx_factory):
-- 
GitLab