def test_copy_statefun_address_into_task_request(self):
address = Address(namespace="tests", type="test", id="id")
task_request = TaskRequest()
task_request.reply_address.ParseFromString(address.SerializeToString())
self.assertEqual(task_request.reply_address.namespace, address.namespace)
self.assertEqual(task_request.reply_address.type, address.type)
self.assertEqual(task_request.reply_address.id, address.id)
def run_pipeline(self, pipeline: PipelineBuilder, initial_target_type='worker'):
task_request = pipeline.to_task_request(serialiser)
task_request.reply_topic = self.__reply_topic
target = Address()
target.namespace = default_namespace
target.type = initial_target_type
target.id = task_request.id
return self._run_flink_loop(task_request, target)
def run_action(self,
pipeline: PipelineBuilder,
action: TaskAction,
initial_target_type='worker'):
task_action = TaskActionRequest(id=pipeline.id,
action=action,
reply_topic=self.__reply_topic)
target = Address()
target.namespace = default_namespace
target.type = initial_target_type
target.id = pipeline.id
return self._run_flink_loop(task_action, target)
def test_convert_dict_to_protofbuf(self):
data = {
'int': 123,
'float': 1.23,
'str': '123',
'list': [1,2,3],
'dict': {
'a': 1,
'b': 2
},
'dict_in_list': [1, {'a': 1}],
'proto': Address(namespace="tests", type="test", id="id")
}
proto = _convert_to_proto(data)
reconsituted_data = _convert_from_proto(proto, known_proto_types=[Address])
self.assertEqual(reconsituted_data['int'], 123)
self.assertEqual(reconsituted_data['float'], 1.23)
self.assertEqual(reconsituted_data['str'], '123')
self.assertEqual(reconsituted_data['list'], [1,2,3])
self.assertEqual(reconsituted_data['dict']['a'], 1)
self.assertEqual(reconsituted_data['dict']['b'], 2)
self.assertEqual(reconsituted_data['dict_in_list'][0], 1)
self.assertEqual(reconsituted_data['dict_in_list'][1]['a'], 1)
self.assertTrue(isinstance(reconsituted_data['proto'], Address))
self.assertEqual(reconsituted_data['proto'].namespace, 'tests')
self.assertEqual(reconsituted_data['proto'].type, 'test')
self.assertEqual(reconsituted_data['proto'].id, 'id')
def test_task_entry_serialisation(self):
serialiser = DefaultSerialiser(known_proto_types=[Address])
args = (1, '2', Address(namespace='test'))
kwargs = {'arg': [1, 2, 3]}
parameters = {'is_fruitful': True}
entry = Task.from_fields('task_id',
'task_type',
args,
kwargs,
**parameters,
is_finally=True)
entry.mark_complete()
entry_proto = entry.to_proto(serialiser)
reconsituted_entry = Task.from_proto(entry_proto).unpack(serialiser)
self.assertEqual(entry_proto.task_entry.request.type_url,
'type.googleapis.com/statefun_tasks.ArgsAndKwargs')
self.assertEqual(reconsituted_entry.task_id, entry.task_id)
self.assertEqual(reconsituted_entry.task_type, entry.task_type)
self.assertEqual(reconsituted_entry.is_fruitful, True)
self.assertEqual(reconsituted_entry.is_finally, True)
self.assertEqual(reconsituted_entry.is_complete(), True)
self.assertEqual(reconsituted_entry.to_tuple(), entry.to_tuple())
def test_group_entry_serialisation(self):
serialiser = DefaultSerialiser(known_proto_types=[Address])
args = (1, '2', Address(namespace='test'))
kwargs = {'arg': [1, 2, 3]}
group_entry = Group(group_id='inner_group_id', max_parallelism=10)
group_entry.add_to_group([
Task.from_fields('inner_task_id_1', 'task_type', args, kwargs),
Task.from_fields('inner_task_id_2', 'task_type', args, kwargs)
])
entry = Group(group_id='group_id')
entry.add_to_group([
group_entry,
Task.from_fields('grouped_task_chain_1_1', 'task_type', args,
kwargs),
Task.from_fields('grouped_task_chain_1_2', 'task_type', args,
kwargs)
])
entry.add_to_group([
Task.from_fields('grouped_task_chain_2_1', 'task_type', args,
kwargs)
])
proto = entry.to_proto(serialiser)
reconsituted_entry = Group.from_proto(proto)
self.assertEqual(str(reconsituted_entry), str(entry))
def run_pipeline(self,
pipeline: PipelineBuilder,
initial_target_type='worker'):
task_request = pipeline.to_task_request(test_harness_serialiser)
task_request.reply_topic = self.__reply_topic
target = Address()
target.namespace = default_namespace
target.type = initial_target_type
target.id = task_request.id
result = self._run_flink_loop(task_request, target)
if isinstance(result, TaskErrorException):
raise result
return result
def run_action(self,
pipeline: PipelineBuilder,
action: TaskAction,
initial_target_type='worker'):
task_action = TaskActionRequest(id=pipeline.id,
action=action,
reply_topic=self.__reply_topic)
target = Address()
target.namespace = default_namespace
target.type = initial_target_type
target.id = pipeline.id
result = self._run_flink_loop(task_action, target)
if isinstance(result, TaskErrorException):
raise result
return result
def test_convert_list_to_protofbuf(self):
data = [Address(namespace="tests", type="test", id="id"), 1, '123']
proto = _convert_to_proto(data)
reconsituted_data = _convert_from_proto(proto,
known_proto_types=[Address])
self.assertTrue(isinstance(reconsituted_data[0], Address))
self.assertEqual(reconsituted_data[0].namespace, 'tests')
self.assertEqual(reconsituted_data[0].type, 'test')
self.assertEqual(reconsituted_data[0].id, 'id')
self.assertEqual(reconsituted_data[1], 1)
self.assertEqual(reconsituted_data[2], '123')
def test_convert_list_to_protofbuf(self):
data = [Address(namespace="tests", type="test", id="id"), 1, '123']
proto = _convert_to_proto(data, _generate_default_converters())
reconsituted_data = _convert_from_proto(proto, [Address],
self.default_converters)
self.assertTrue(isinstance(reconsituted_data[0], Address))
self.assertEqual(reconsituted_data[0].namespace, 'tests')
self.assertEqual(reconsituted_data[0].type, 'test')
self.assertEqual(reconsituted_data[0].id, 'id')
self.assertEqual(reconsituted_data[1], 1)
self.assertEqual(reconsituted_data[2], '123')
def test_task_entry_serialisation_with_single_protobuf_arg(self):
serialiser = DefaultSerialiser(known_proto_types=[Address])
args = Address(namespace='test')
entry = Task.from_fields('task_id', 'task_type', args, {}, True)
entry_proto = entry.to_proto(serialiser)
reconsituted_entry = Task.from_proto(entry_proto).unpack(serialiser)
self.assertEqual(
entry_proto.task_entry.request.type_url,
'type.googleapis.com/io.statefun.sdk.reqreply.Address')
self.assertEqual(reconsituted_entry.to_tuple(), entry.to_tuple())
def test_task_entry_serialisation_with_single_protobuf_arg(self):
serialiser = DefaultSerialiser(known_proto_types=[Address])
args = Address(namespace='test')
entry = _TaskEntry('task_id', 'task_type', args, {}, {}, True)
entry_proto = entry.to_proto(serialiser)
reconsituted_entry = _TaskEntry.from_proto(entry_proto, serialiser)
self.assertEqual(
entry_proto.task_entry.request.type_url,
'type.googleapis.com/org.apache.flink.statefun.flink.core.polyglot.Address'
)
self.assertEqual(reconsituted_entry.args, entry.args)
self.assertEqual(reconsituted_entry.kwargs, {})
def test_task_entry_serialisation(self):
serialiser = DefaultSerialiser(known_proto_types=[Address])
args = (1, '2', Address(namespace='test'))
kwargs = {'arg': [1, 2, 3]}
parameters = {'a_parameter': 'some_value'}
entry = _TaskEntry('task_id', 'task_type', args, kwargs, parameters,
True)
entry.mark_complete()
entry_proto = entry.to_proto(serialiser)
reconsituted_entry = _TaskEntry.from_proto(entry_proto, serialiser)
self.assertEqual(reconsituted_entry.task_id, entry.task_id)
self.assertEqual(reconsituted_entry.task_type, entry.task_type)
self.assertEqual(reconsituted_entry.args, tuple(entry.args, ))
self.assertEqual(reconsituted_entry.kwargs, kwargs)
self.assertEqual(reconsituted_entry.parameters, parameters)
self.assertEqual(reconsituted_entry.is_finally, True)
self.assertEqual(reconsituted_entry.is_complete(), True)
def update_address(address: Address, namespace, address_type, address_id):
address.namespace = namespace
address.type = address_type
address.id = address_id
