def test_pipeline_with_group_that_waits(self): pipeline = in_parallel([ tasks.send(_say_hello, 'Jane', 'Doe'), tasks.send(_say_hello, 'Joe', 'Blogs') ]).wait().continue_with(_say_goodbye, goodbye_message="see you later!") self.test_harness.run_pipeline(pipeline) entries = pipeline.get_tasks() self.assertIn(entries[0][2], _started) # first task in group self.assertIn(entries[1][2], _started) # second task in group self.assertNotIn(entries[2][2], _started) # task after group self.assertIn(pipeline.id, _status) self.assertEqual(_status[pipeline.id], TaskStatus.PAUSED)
def test_get_task_request_for_an_existing_pipeline(self): pipeline = tasks.send(_say_hello, 'Jane', 'Doe') self.test_harness.run_pipeline(pipeline) action_result = self.test_harness.run_action(pipeline, TaskAction.GET_REQUEST) task_request = self._unpack(action_result.result, TaskRequest) self.assertEqual(task_request.id, pipeline.id)
def test_pipeline_with_continuation(self): pipeline = tasks.send(hello_and_goodbye_workflow, 'Jane', last_name='Doe') result = self.test_harness.run_pipeline(pipeline) self.assertEqual(result, 'Hello Jane Doe. So now I will say see you later!')
def test_get_status_for_pending_pipeline(self): pipeline = tasks.send(_say_hello, 'Jane', 'Doe') action_result = self.test_harness.run_action(pipeline, TaskAction.GET_STATUS) self.assertEqual( self._unpack(action_result.result, TaskStatus).status, TaskStatus.PENDING)
def test_get_task_result_for_a_completed_pipeline(self): pipeline = tasks.send(_say_hello, 'Jane', 'Doe') self.test_harness.run_pipeline(pipeline) action_result = self.test_harness.run_action(pipeline, TaskAction.GET_RESULT) task_result = self._unpack(action_result.result, TaskResult) self.assertEqual(task_result.id, pipeline.id)
def test_providing_mandatory_plus_one_default_arg(self): pipeline = tasks.send(multi_arg_workflow_with_some_defaults, 'a', 'b', d='d') result = self.test_harness.run_pipeline(pipeline) self.assertEqual(result, 'a,b,c,d')
def test_pipeline_with_finally_state(self): global finally_state finally_state = 0 pipeline = tasks.send(workflow_with_cleanup_state, initial_state=10) self.test_harness.run_pipeline(pipeline) self.assertEqual(finally_state, 20) # finally adds 10
def test_pipeline_using_kwargs(self): pipeline = tasks.send(hello_workflow, first_name='Jane', last_name='Doe') proto = pipeline.to_proto(serialiser=DefaultSerialiser()) self.assertEqual(proto.entries[0].task_entry.request.type_url, 'type.googleapis.com/statefun_tasks.ArgsAndKwargs') result = self.test_harness.run_pipeline(pipeline) self.assertEqual(result, 'Hello Jane Doe')
def test_passing_state_through_exception_throwing_task(self): pipeline = tasks.send(state_passing_workflow_with_exception) with self.assertRaises(Exception): self.test_harness.run_pipeline(pipeline) self.assertEqual(_final_state, 4560)
def hello_and_goodbye_workflow(context, first_name, last_name): if not (context.get_root_pipeline_id() == _root_pipeline_id): raise ValueError('Root ID Mismatch') return tasks.send(_say_hello, first_name, last_name).continue_with( _say_goodbye, goodbye_message="see you later!")
def test_get_status_for_completed_pipeline(self): pipeline = tasks.send(_say_hello, 'Jane', 'Doe') self.test_harness.run_pipeline(pipeline) action_result = self.test_harness.run_action(pipeline, TaskAction.GET_STATUS) self.assertEqual( self._unpack(action_result.result, TaskStatus).status, TaskStatus.COMPLETED)
def test_failing_pipeline_4_times_with_3_retries(self): pipeline = tasks.send(fail_workflow, 4) try: self.test_harness.run_pipeline(pipeline) except TaskErrorException as e: self.assertEqual(e.task_error.message, 'Failed after 4 iterations') else: self.fail('Expected an exception')
def test_sending_protobuf_to_function_without_annotations_packs_as_any(self): pipeline = tasks.send(receive_and_reply_protobuf_not_annotated, UnknownProto()) result = self.test_harness.run_pipeline(pipeline) self.assertIsInstance(result, Any) unknown = UnknownProto() result.Unpack(unknown) self.assertEqual("<class 'google.protobuf.any_pb2.Any'>", unknown.value_str)
def test_sending_request_data_that_is_not_protobuf_serialisable(self): pipeline = tasks.send(return_my_class_my_field, MyClass()) try: self.test_harness.run_pipeline(pipeline) except Exception as e: self.assertIn('Cannot convert value', str(e)) else: self.fail('Expected an exception')
def test_get_task_request_for_a_non_existing_pipeline(self): pipeline = tasks.send(_say_hello, 'Jane', 'Doe') try: self.test_harness.run_action(pipeline, TaskAction.GET_REQUEST) except TaskErrorException as ex: self.assertEqual(ex.task_error.message, 'Task request not found') else: self.fail('Expected an exception')
def test_sending_protobuf_to_function_with_annotations_uses_exact_protos( self): pipeline = tasks.send(receive_and_reply_protobuf_fully_annotated, TestProto()) result = self.test_harness.run_pipeline(pipeline) self.assertIsInstance(result, TestResultProto) self.assertEqual("<class 'test_messages_pb2.TestProto'>", result.value_str)
def test_simple_protobuf_pipeline(self): pipeline = tasks.send(simple_protobuf_workflow, TestPerson(first_name='Jane', last_name='Doe')) proto = pipeline.to_proto(serialiser=DefaultSerialiser()) self.assertEqual(proto.entries[0].task_entry.request.type_url, 'type.googleapis.com/tests.TestPerson') result = self.test_harness.run_pipeline(pipeline) self.assertEqual(result.greeting, 'Hello Jane Doe')
def hello_workflow(context, first_name, last_name): global _root_pipeline_id if not (context.get_root_pipeline_id() == context.get_pipeline_id()): raise ValueError('Root ID Mismatch') _root_pipeline_id = context.get_root_pipeline_id() return tasks.send(hello_and_goodbye_workflow, first_name, last_name)
def test_pipeline_with_unknown_worker_name(self): pipeline = tasks.send(hello_workflow, 'Jane', 'Doe').set(worker_name='unknown') try: self.test_harness.run_pipeline(pipeline) except KeyError as e: self.assertEqual(str(e.args[0]), 'test/unknown') else: self.fail('Expected an exception')
def test_pipeline_with_finally_kwargs(self): global finally_args finally_args = None pipeline = tasks.send(workflow_with_cleanup_kwargs, 'Jane', last_name='Doe') self.test_harness.run_pipeline(pipeline) self.assertEqual(finally_args, ['arg1', 'arg2'])
def test_pipeline_with_finally(self): global finally_flag finally_flag = False pipeline = tasks.send(workflow, 'Jane', last_name='Doe') result = self.test_harness.run_pipeline(pipeline) self.assertEqual(result, 'Hello Jane Doe. So now I will say see you later!') self.assertEqual(finally_flag, True)
def test_pipeline_tasks_start_and_end(self): pipeline = tasks.send(hello_workflow, 'Jane', 'Doe') self.test_harness.run_pipeline(pipeline) self.assertIn(pipeline.id, _started) self.assertIn(pipeline.id, _finished) for _, _, task_id in pipeline.get_tasks(): self.assertIn(task_id, _started) self.assertIn(task_id, _finished) self.assertIn(task_id, _pipeline_tasks_finished)
def test_sending_primitives_to_function(self): pipeline = tasks.send(receive_and_reply_primitives, 1, 2.0, [3], 'four') result = self.test_harness.run_pipeline(pipeline) self.assertIsInstance(result, list) a, b, c, d = result self.assertEqual(1, a) self.assertEqual(2.0, b) self.assertEqual({'c': [3]}, c) self.assertEqual(('four', ), d)
def test_get_tassk_result_for_failed_pipeline(self): pipeline = tasks.send(_say_hello, 'Jane') try: self.test_harness.run_pipeline(pipeline) except: pass action_result = self.test_harness.run_action(pipeline, TaskAction.GET_RESULT) task_exception = self._unpack(action_result.result, TaskException) self.assertEqual(task_exception.id, pipeline.id)
def test_not_providing_mandatory_args(self): pipeline = tasks.send(multi_arg_workflow_with_some_defaults, 'a', c='c', d='d') try: self.test_harness.run_pipeline(pipeline) except TaskErrorException as e: self.assertEqual(e.task_error.message, 'Not enough args supplied') else: self.fail('Expected an exception')
def test_get_status_for_failed_pipeline(self): pipeline = tasks.send(_say_hello, 'Jane') try: self.test_harness.run_pipeline(pipeline) except: pass action_result = self.test_harness.run_action(pipeline, TaskAction.GET_STATUS) self.assertEqual( self._unpack(action_result.result, TaskStatus).status, TaskStatus.FAILED)
def test_sending_embedded_protobuf_in_primitives_to_function(self): pipeline = tasks.send(receive_and_reply_primitives, 1, TestProto(), [3], UnknownProto()) result = self.test_harness.run_pipeline(pipeline) self.assertIsInstance(result, list) a, b, c, d = result self.assertEqual(1, a) self.assertIsInstance(b, TestProto) self.assertEqual({'c': [3]}, c) p, = d self.assertIsInstance(p, Any)
def test_unregistered_function(self): @other_tasks_instance.bind() def new_function(): return 'Hello' pipeline = tasks.send(new_function, 'Jane', 'Doe') try: self.test_harness.run_pipeline(pipeline) except TaskErrorException as e: self.assertEqual( str(e.task_error.message), f'{__name__}.new_function is not a registered FlinkTask') else: self.fail('Expected a TaskErrorException')
def test_pipeline_with_cleanup_in_middle(self): pipeline = tasks.send(workflow_with_cleanup_in_middle, 'Jane', last_name='Doe') try: self.test_harness.run_pipeline(pipeline) except TaskErrorException as e: self.assertEqual( e.task_error.message, 'Invalid pipeline: "finally_do" must be called at the end of a pipeline' ) else: self.fail('Expected an exception')
def test_pipeline_with_finally_and_error(self): global finally_flag finally_flag = False pipeline = tasks.send(workflow_with_error_and_cleanup, 'Jane', last_name='Doe') try: self.test_harness.run_pipeline(pipeline) except TaskErrorException as e: self.assertIn('I am supposed to fail', e.task_error.message) self.assertEqual(finally_flag, True) else: self.fail('Expected an exception')