ray/streaming/python/tests/test_operator.py
2020-06-08 14:28:11 +08:00

37 lines
1.3 KiB
Python

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_with_func():
map_func = function.SimpleMapFunction(lambda x: x)
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)