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')
