ray/streaming/python/tests/test_function.py
2020-02-25 10:33:33 +08:00

22 lines
697 B
Python

from ray.streaming import function
from ray.streaming.runtime import gateway_client
def test_get_simple_function_class():
simple_map_func_class = function._get_simple_function_class(
function.MapFunction)
assert simple_map_func_class is function.SimpleMapFunction
class MapFunc(function.MapFunction):
def map(self, value):
return str(value)
def test_load_function():
# function_bytes, module_name, class_name, function_name,
# function_interface
descriptor_func_bytes = gateway_client.serialize(
[None, __name__, MapFunc.__name__, None, "MapFunction"])
func = function.load_function(descriptor_func_bytes)
assert type(func) is MapFunc