A few fixes in receive() signal. (#4142)

python/ray/tests/test_signal.py

@@ -1,4 +1,5 @@
 import pytest
+import time
 
 import ray
 import ray.experimental.signal as signal
@@ -205,8 +206,10 @@ def test_actor_crash_init3(ray_start):
 
     a = ActorCrashInit.remote()
     a.method.remote()
-    result_list = signal.receive([a], timeout=10)
-    assert len(result_list) == 1
+    # Wait for a.method.remote() to finish and generate an error.
+    time.sleep(10)
+    result_list = signal.receive([a], timeout=5)
+    assert len(result_list) == 2
     assert type(result_list[0][1]) == signal.ErrorSignal
 
 
@@ -279,3 +282,37 @@ def test_forget(ray_start):
     ray.get(a.send_signals.remote(signal_value, count))
     result_list = receive_all_signals([a], timeout=5)
     assert len(result_list) == count
+
+
+def test_send_signal_from_two_tasks_to_driver(ray_start):
+    # Define a remote function that sends a user-defined signal.
+    @ray.remote
+    def send_signal(value):
+        signal.send(UserSignal(value))
+
+    a = send_signal.remote(0)
+    b = send_signal.remote(0)
+
+    ray.get([a, b])
+
+    result_list = ray.experimental.signal.receive([a])
+    assert len(result_list) == 1
+    # Call again receive on "a" with no new signal.
+    result_list = ray.experimental.signal.receive([a, b])
+    assert len(result_list) == 1
+
+
+def test_receiving_on_two_returns(ray_start):
+    @ray.remote(num_return_vals=2)
+    def send_signal(value):
+        signal.send(UserSignal(value))
+        return 1, 2
+
+    x, y = send_signal.remote(0)
+
+    ray.get([x, y])
+
+    results = ray.experimental.signal.receive([x, y])
+
+    assert ((x == results[0][0] and y == results[1][0])
+            or (x == results[1][0] and y == results[0][0]))