def test_task_io(self):
try:
flow = workflow.Workflow()
task1 = Task1(workflow=flow, task_id='task1')
task2 = Task2(workflow=flow, task_id='task2')
task2.bind_input_to_task_output('task1', task1)
flow.build()
flow.process(nworkers=10)
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.SUCCESS, task2.state)
flow2 = workflow.Workflow()
task1_again = Task1(workflow=flow2, task_id='task1')
task2_again = Task2(workflow=flow2, task_id='task2')
task2_again.bind_input_to_task_output('task1', task1_again)
task3 = Task1(workflow=flow2, task_id='task3', runs_after=[task2_again])
flow2.build()
flow2.process(nworkers=10)
for original, again in ((task1, task1_again), (task2, task2_again)):
self.assertEqual(workflow.TaskState.ALREADY_DONE, again.state)
self.assertEqual(original.start_time, again.start_time)
self.assertEqual(original.end_time, again.end_time)
self.assertEqual(workflow.TaskState.SUCCESS, task3.state)
finally:
os.remove(task1.get_task_run_id())
os.remove(task2.get_task_run_id())
os.remove(task3.get_task_run_id())
def test_workflow_task_error_calls_shutdown(self):
real_shutdown = base.shutdown
try:
shutdown_count = 0
flow = workflow.Workflow()
def mock_shutdown():
nonlocal shutdown_count
nonlocal flow
logging.info("base.shutdown called")
shutdown_count += 1
flow._done.set()
base.shutdown = mock_shutdown
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = ErrorTask(workflow=flow, task_id="task2", runs_after=[task1])
task3 = SuccessTask(workflow=flow, task_id="task3", runs_after=[task2])
flow.build()
flow.process(nworkers=10)
self.assertTrue(shutdown_count > 0, "shutdown should have been called on error")
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
# task2.state is undefined because base.shutdown didn't shutdown
self.assertEqual(workflow.TaskState.PENDING, task3.state)
finally:
base.shutdown = real_shutdown
def test_dump_state_as_table(self):
flow1 = workflow.Workflow()
task1 = SlowTask(workflow=flow1, task_id='task1')
task2 = SlowTask(workflow=flow1, task_id='task2')
task3 = SlowTask(workflow=flow1, task_id='task3')
task2.must_run_after(task1)
task3.must_run_after(task2)
flow1.build()
svg_source = flow1.dump_state_as_table()
logging.debug('Workflow table:\n%s', svg_source)
server = workflow.WorkflowHTTPMonitor(interface='127.0.0.1', port=0)
server.start()
server.set_workflow(flow1)
try:
flow1.process()
conn = http.client.HTTPConnection(
host='127.0.0.1',
port=server.server_port,
)
conn.connect()
try:
conn.request(method='GET', url='')
response = conn.getresponse()
self.assertEqual(200, response.getcode())
self.assertEqual('text/plain',
response.getheader('Content-Type'))
logging.debug('HTTP response: %r', response.read())
finally:
conn.close()
finally:
server.stop()
def test_completion_handler_exceptions(self):
"""Tests a failed completion handler does not prevent others from running."""
completion_handler_workflow_pointer = None
def failing_completion_handler(workflow):
"""Function to run after the workflow completes.
Args:
workflow: The workflow that completed."""
nonlocal completion_handler_workflow_pointer
completion_handler_workflow_pointer = workflow
raise Error("This completion handler is supposed to error.")
completion_handler = mock.Mock()
flow = workflow.Workflow()
flow.add_completion_handler(failing_completion_handler)
flow.add_completion_handler(completion_handler)
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = FailureTask(workflow=flow, task_id='task2', runs_after=[task1])
flow.build()
flow.process(nworkers=10)
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.FAILURE, task2.state)
self.assertIs(flow, completion_handler_workflow_pointer)
self.assertEqual(1, completion_handler.call_count)
def testAddDep(self):
"""Tests that forward declaration of dependencies are properly set."""
flow = workflow.Workflow()
flow.AddDep('task1', 'task2')
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = SuccessTask(workflow=flow, task_id='task2')
self.assertTrue('task1' in task2.runs_after)
self.assertTrue('task2' in task1.runs_before)
def disabled_test_workflow_diff(self):
flow1 = workflow.Workflow()
task1 = SuccessTask(workflow=flow1, task_id='task1')
task2 = SuccessTask(workflow=flow1, task_id='task2')
task3 = SuccessTask(workflow=flow1, task_id='task3')
task2.must_run_after(task1)
task3.must_run_after(task2)
flow1.build()
flow2 = workflow.Workflow()
task1 = SuccessTask(workflow=flow2, task_id='task1')
task2b = SuccessTask(workflow=flow2, task_id='task2b')
task3 = SuccessTask(workflow=flow2, task_id='task3')
task2b.must_run_after(task1)
task3.must_run_after(task2b)
task3.must_run_after(task1)
flow2.build()
workflow.diff_workflow(flow1, flow2)
def test_should_task_run(self):
class TestTask(workflow.IOTask):
def __init__(self, should_task_run=True, **kwargs):
super().__init__(**kwargs)
self._should_counter = 0
self._run_counter = 0
self._should_task_run = should_task_run
def should_task_run(self, task_run_id, output, **inputs):
self._should_counter += 1
return self._should_task_run
def run_with_io(self, output, **inputs):
self._run_counter += 1
output.time = time.time()
return self.SUCCESS
flow = workflow.Workflow()
task = TestTask(workflow=flow, task_id="TestTask")
flow.build()
flow.process()
self.assertEqual(0, task._should_counter)
self.assertEqual(1, task._run_counter)
flow = workflow.Workflow()
task = TestTask(workflow=flow, task_id="TestTask", should_task_run=True)
flow.build()
flow.process()
self.assertEqual(1, task._should_counter)
self.assertEqual(1, task._run_counter)
flow = workflow.Workflow()
task = TestTask(workflow=flow, task_id="TestTask", should_task_run=False)
flow.build()
flow.process()
self.assertEqual(1, task._should_counter)
self.assertEqual(0, task._run_counter)
os.remove(task.get_task_run_id())
def testCircularDep(self):
flow = workflow.Workflow()
task1 = SuccessTask(workflow=flow, task_id='task1', runs_after=['task3'])
task2 = FailureTask(workflow=flow, task_id='task2', runs_after=[task1])
task3 = SuccessTask(workflow=flow, task_id='task3', runs_after=[task2])
try:
flow.build()
self.fail("circular dependency task1->%s, task2->%s, task3->%s"
% (task1.runs_after, task2.runs_after, task3.runs_after))
except workflow.CircularDependencyError:
pass
def testCircularDep2(self):
flow = workflow.Workflow()
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = FailureTask(workflow=flow, task_id='task2')
task3 = SuccessTask(workflow=flow, task_id='task3')
task2.must_run_after(task3)
task2.must_run_before(task3)
try:
flow.build()
self.fail()
except workflow.CircularDependencyError:
pass
def testWorkflowFailure(self):
flow = workflow.Workflow()
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = FailureTask(workflow=flow, task_id='task2', runs_after=[task1])
task3 = SuccessTask(workflow=flow, task_id='task3', runs_after=[task2])
task4 = SuccessTask(workflow=flow, task_id='task4', runs_after=[task2, task3])
flow.build()
flow.process(nworkers=10)
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.FAILURE, task2.state)
self.assertEqual(workflow.TaskState.FAILURE, task3.state)
self.assertEqual(workflow.TaskState.FAILURE, task4.state)
def test_get_downstream_tasks(self):
flow = workflow.Workflow()
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = SuccessTask(workflow=flow, task_id='task2')
task3 = SuccessTask(workflow=flow, task_id='task3')
task4 = SuccessTask(workflow=flow, task_id='task4')
task5 = SuccessTask(workflow=flow, task_id='task5')
task6 = SuccessTask(workflow=flow, task_id='task6')
task7 = SuccessTask(workflow=flow, task_id='task7')
task3.must_run_after(task1)
task3.must_run_after(task2)
task4.must_run_after(task3)
task5.must_run_after(task3)
task7.must_run_after(task6)
self.assertEqual(
set({task1, task3, task4, task5}),
workflow.get_downstream_tasks(flow, [task1]))
self.assertEqual(
set({task3, task4, task5}),
workflow.get_downstream_tasks(flow, [task3]))
self.assertEqual(
set({task4}),
workflow.get_downstream_tasks(flow, [task4]))
self.assertEqual(
set({task5}),
workflow.get_downstream_tasks(flow, [task5]))
self.assertEqual(
set({task4, task5}),
workflow.get_downstream_tasks(flow, [task4, task5]))
self.assertEqual(
set({task3, task4, task5}),
workflow.get_downstream_tasks(flow, [task3, task4, task5]))
self.assertEqual(
set({task6, task7}),
workflow.get_downstream_tasks(flow, [task6]))
self.assertEqual(
set({task7}),
workflow.get_downstream_tasks(flow, [task7]))
self.assertEqual(
set({task3, task4, task5, task6, task7}),
workflow.get_downstream_tasks(flow, [task3, task6]))
def testParallelWorkflowSuccess(self):
flow = workflow.Workflow()
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = SuccessTask(workflow=flow, task_id='task2')
task3 = SuccessTask(workflow=flow, task_id='task3')
task4 = SuccessTask(workflow=flow, task_id='task4')
task5 = SuccessTask(workflow=flow, task_id='task5')
flow.build()
flow.process(nworkers=10)
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.SUCCESS, task2.state)
self.assertEqual(workflow.TaskState.SUCCESS, task3.state)
self.assertEqual(workflow.TaskState.SUCCESS, task4.state)
self.assertEqual(workflow.TaskState.SUCCESS, task5.state)
def testLinearWorkflowSuccess(self):
flow = workflow.Workflow()
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = SuccessTask(workflow=flow, task_id='task2', runs_after=[task1])
task3 = SuccessTask(workflow=flow, task_id='task3', runs_after=[task2])
task4 = SuccessTask(workflow=flow, task_id='task4', runs_after=[task3])
task5 = SuccessTask(workflow=flow, task_id='task5', runs_after=[task4])
flow.build()
flow.process(nworkers=10)
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.SUCCESS, task2.state)
self.assertEqual(workflow.TaskState.SUCCESS, task3.state)
self.assertEqual(workflow.TaskState.SUCCESS, task4.state)
self.assertEqual(workflow.TaskState.SUCCESS, task5.state)
print(flow.dump_as_dot())
def test_completion_handlers_multiple(self):
"""Tests that multiple completion handlers run after a successful workflow run."""
completion_handler_1 = mock.Mock()
completion_handler_2 = mock.Mock()
flow = workflow.Workflow()
flow.add_completion_handler(completion_handler_1)
flow.add_completion_handler(completion_handler_2)
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = SuccessTask(workflow=flow, task_id='task2', runs_after=[task1])
flow.build()
flow.process(nworkers=3)
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.SUCCESS, task2.state)
self.assertEqual(1, completion_handler_1.call_count)
self.assertEqual(1, completion_handler_2.call_count)
def test_completion_handler_empty_workflow(self):
"""Tests that the completion handler is run after an empty workflow run."""
completion_handler_workflow_pointer = None
def test_completion_handler(workflow):
"""Function to run after the workflow completes.
Args:
workflow: The workflow that completed."""
nonlocal completion_handler_workflow_pointer
completion_handler_workflow_pointer = workflow
flow = workflow.Workflow()
flow.add_completion_handler(test_completion_handler)
flow.build()
flow.process(nworkers=3)
self.assertIs(flow, completion_handler_workflow_pointer)
def testPersistentTaskRunCreateFile(self):
"""Running a persistent task with success creates the output file."""
class PTask(workflow.LocalFSPersistentTask):
def run(self):
return self.SUCCESS
flow = workflow.Workflow()
file_path = base.random_alpha_num_word(16)
try:
task = PTask(
output_file_path=file_path,
task_id='task',
workflow=flow,
)
flow.build()
self.assertFalse(os.path.exists(file_path))
self.assertTrue(flow.process())
self.assertTrue(os.path.exists(file_path))
finally:
base.remove(file_path)
def testPersistentTaskRunNoFile(self):
"""Persistent task runs if output file does not exist.
No output file is created if task fails.
"""
class PTask(workflow.LocalFSPersistentTask):
def run(self):
return self.FAILURE
flow = workflow.Workflow()
file_path = base.random_alpha_num_word(16)
task = PTask(
output_file_path=file_path,
task_id='task',
workflow=flow,
)
flow.build()
self.assertFalse(os.path.exists(file_path))
self.assertFalse(flow.process())
self.assertFalse(os.path.exists(file_path))
def test_prune(self):
"""Tests that forward declaration of dependencies are properly set."""
flow = workflow.Workflow()
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = SuccessTask(workflow=flow, task_id='task2')
task3 = SuccessTask(workflow=flow, task_id='task3')
task4 = SuccessTask(workflow=flow, task_id='task4')
task5 = SuccessTask(workflow=flow, task_id='task5')
task6 = SuccessTask(workflow=flow, task_id='task6')
task7 = SuccessTask(workflow=flow, task_id='task7')
task3.must_run_after(task1)
task3.must_run_after(task2)
task4.must_run_after(task3)
task5.must_run_after(task3)
task7.must_run_after(task6)
flow.prune(tasks={task4}, direction=workflow.UPSTREAM)
self.assertEqual({'task1', 'task2', 'task3', 'task4'}, flow.tasks.keys())
def testPersistentTaskNoRunIfFile(self):
"""Persistent task does not run if output file already exists."""
class PTask(workflow.LocalFSPersistentTask):
def run(self):
return self.FAILURE
flow = workflow.Workflow()
file_path = base.random_alpha_num_word(16)
base.touch(file_path)
try:
task = PTask(
output_file_path=file_path,
task_id='task',
workflow=flow,
)
flow.build()
self.assertTrue(os.path.exists(file_path))
self.assertTrue(flow.process())
self.assertTrue(os.path.exists(file_path))
finally:
base.remove(file_path)
def test_completion_handler_on_failure(self):
"""Tests that the completion handler is run after an unsuccessful workflow run."""
completion_handler_workflow_pointer = None
def test_completion_handler(workflow):
"""Function to run after the workflow completes.
Args:
workflow: The workflow that completed."""
nonlocal completion_handler_workflow_pointer
completion_handler_workflow_pointer = workflow
flow = workflow.Workflow()
flow.add_completion_handler(test_completion_handler)
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = FailureTask(workflow=flow, task_id='task2', runs_after=[task1])
flow.build()
flow.process(nworkers=10)
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.FAILURE, task2.state)
self.assertIs(flow, completion_handler_workflow_pointer)
def test_workflow_task_error_handler(self):
task_error_kwargs = None
def handle_task_error(**kwargs):
nonlocal task_error_kwargs
task_error_kwargs = kwargs
return workflow.TaskState.SUCCESS
flow = workflow.Workflow(task_error_handler=handle_task_error)
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = ErrorTask(workflow=flow, task_id="task2", runs_after=[task1])
task3 = SuccessTask(workflow=flow, task_id="task3", runs_after=[task2])
flow.build()
flow.process(nworkers=10)
self.assertEqual(flow, task_error_kwargs["flow"])
self.assertEqual(task2, task_error_kwargs["task"])
self.assertEqual(workflow.Error, type(task_error_kwargs["error"]))
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.SUCCESS, task2.state)
self.assertEqual(workflow.TaskState.SUCCESS, task3.state)
def test_workflow_async(self):
"""Tests that a workflow can run successfully asynchronously, and that the completion
handler also runs."""
completion_handler_workflow_pointer = None
def test_completion_handler(workflow):
"""Function to run after the workflow completes.
Args:
workflow: The workflow that completed."""
nonlocal completion_handler_workflow_pointer
completion_handler_workflow_pointer = workflow
flow = workflow.Workflow()
flow.add_completion_handler(completion_handler=test_completion_handler)
task1 = SuccessTask(workflow=flow, task_id='task1')
task2 = SuccessTask(workflow=flow, task_id='task2', runs_after=[task1])
flow.build()
flow.process(nworkers=3, sync=False)
flow.wait()
self.assertEqual(workflow.TaskState.SUCCESS, task1.state)
self.assertEqual(workflow.TaskState.SUCCESS, task2.state)
self.assertIs(flow, completion_handler_workflow_pointer)
def test_empty_workflow(self):
flow = workflow.Workflow()
flow.build()
flow.process()
def _make_workflow(self, config, force_build=False):
"""Builds the workflow tasks after the build graph defined by the workspace.
Args:
config: Workspace configuration record.
force_build: When true, force a rebuild irrespective of incremental changes.
Returns:
The workflow builder.
"""
workflow_name = "build.commit={commit}.timestamp={timestamp}".format(
commit=self.git.get_commit_hash(),
timestamp=base.timestamp(),
)
flow = workflow.Workflow(name=workflow_name)
TASK_CLASS_MAP = dict(
avro_java_library=workflow_task.AvroJavaLibraryTask,
generated_pom=workflow_task.GeneratePomTask,
java_binary=workflow_task.JavaBinaryTask,
java_library=workflow_task.JavaLibraryTask,
java_test=workflow_task.JavaTestTask,
java_super_binary=workflow_task.JavaSuperBinaryTask,
js_app=workflow_task.JSAppTask,
npm_install=workflow_task.NpmInstallTask,
bower_install=workflow_task.BowerInstallTask,
python_binary=workflow_task.PythonBinaryTask,
python_library=workflow_task.PythonLibraryTask,
python_test=workflow_task.PythonTestTask,
run_checkstyle=workflow_task.RunCheckstyleTask,
run_java_test=workflow_task.RunJavaTestTask,
run_python_test=workflow_task.RunPythonTestTask,
run_scala_test=workflow_task.RunJavaTestTask,
run_scalastyle=workflow_task.RunScalastyleTask,
scala_library=workflow_task.ScalaLibraryTask,
scala_test=workflow_task.ScalaTestTask,
)
for name, definition in self._build_defs.definitions.items():
task_class = TASK_CLASS_MAP[definition.kind]
task = task_class(
workflow=flow,
workspace=self,
name=name,
spec=definition,
force=force_build,
)
for dep in definition.get("deps", tuple()):
if isinstance(dep, artifact.Artifact):
# Direct Maven artifact dependencies are not reified as workflow tasks:
continue
if isinstance(dep, build_defs.PythonPyPIDep):
# Python PyPI dependencies are external dependencies:
continue
if isinstance(dep, build_defs.DynamicDep):
dep = dep.provider
if dep is None:
continue
task.bind_input_to_task_output(input_name=dep, task=dep)
return flow
def setUp(self):
self.flow = workflow.Workflow(name="task_test")
self.run_order = []
