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