[Streaming] Union api (#8612)

streaming/python/tests/test_operator.py

@@ -1,8 +1,37 @@
-from ray.streaming import operator
 from ray.streaming import function
+from ray.streaming import operator
+from ray.streaming.operator import OperatorType
+from ray.streaming.runtime import gateway_client
 
 
-def test_create_operator():
+def test_create_operator_with_func():
     map_func = function.SimpleMapFunction(lambda x: x)
-    map_operator = operator.create_operator(map_func)
+    map_operator = operator.create_operator_with_func(map_func)
     assert type(map_operator) is operator.MapOperator
+
+
+class MapFunc(function.MapFunction):
+    def map(self, value):
+        return str(value)
+
+
+class EmptyOperator(operator.StreamOperator):
+    def __init__(self):
+        super().__init__(function.EmptyFunction())
+
+    def operator_type(self) -> OperatorType:
+        return OperatorType.ONE_INPUT
+
+
+def test_load_operator():
+    # function_bytes, module_name, class_name,
+    descriptor_func_bytes = gateway_client.serialize(
+        [None, __name__, MapFunc.__name__, "MapFunction"])
+    descriptor_op_bytes = gateway_client.serialize(
+        [descriptor_func_bytes, "", ""])
+    map_operator = operator.load_operator(descriptor_op_bytes)
+    assert type(map_operator) is operator.MapOperator
+    descriptor_op_bytes = gateway_client.serialize(
+        [None, __name__, EmptyOperator.__name__])
+    test_operator = operator.load_operator(descriptor_op_bytes)
+    assert isinstance(test_operator, EmptyOperator)

streaming/python/tests/test_union_stream.py

@@ -0,0 +1,47 @@
+import os
+
+import ray
+from ray.streaming import StreamingContext
+
+
+def test_union_stream():
+    ray.init(load_code_from_local=True, include_java=True)
+    ctx = StreamingContext.Builder() \
+        .option("streaming.metrics.reporters", "") \
+        .build()
+    sink_file = "/tmp/test_union_stream.txt"
+    if os.path.exists(sink_file):
+        os.remove(sink_file)
+
+    def sink_func(x):
+        with open(sink_file, "a") as f:
+            print("sink_func", x)
+            f.write(str(x))
+
+    stream1 = ctx.from_values(1, 2)
+    stream2 = ctx.from_values(3, 4)
+    stream3 = ctx.from_values(5, 6)
+    stream1.union(stream2, stream3).sink(sink_func)
+    ctx.submit("test_union_stream")
+    import time
+    slept_time = 0
+    while True:
+        if os.path.exists(sink_file):
+            time.sleep(3)
+            with open(sink_file, "r") as f:
+                result = f.read()
+                print("sink result", result)
+                assert set(result) == {"1", "2", "3", "4", "5", "6"}
+            print("Execution succeed")
+            break
+        if slept_time >= 60:
+            raise Exception("Execution not finished")
+        slept_time = slept_time + 1
+        print("Wait finish...")
+        time.sleep(1)
+
+    ray.shutdown()
+
+
+if __name__ == "__main__":
+    test_union_stream()