def test_convert_relations(self):
op1 = TaskGroup(
name="task1",
tasks=[self._create_task("task1")],
downstream_names=["task2", "task3"],
error_downstream_name="fail1",
)
op2 = TaskGroup(
name="task2",
tasks=[self._create_task("task2")],
downstream_names=["task3", "task4"],
error_downstream_name="fail1",
)
op3 = TaskGroup(
name="task3",
tasks=[self._create_task("task3")],
downstream_names=["end1"],
error_downstream_name="fail1",
)
op4 = mock.Mock(
**{
"first_task_id": "task4_first",
"last_task_id_of_ok_flow": "task4_last",
"last_task_id_of_error_flow": "task4_last",
"downstream_names": ["task1", "task2", "task3"],
"error_downstream_name": "fail1",
}
)
end = TaskGroup(name="end1", tasks=[self._create_task("end1")])
fail = TaskGroup(name="fail1", tasks=[self._create_task("fail1")])
op_dict = {"task1": op1, "task2": op2, "task3": op3, "task4": op4, "end1": end, "fail1": fail}
workflow = self._create_workflow(task_groups=op_dict)
converter = self._create_converter()
workflow.task_group_relations = set()
converter.workflow = workflow
converter.add_state_handlers()
converter.convert_relations()
self.assertEqual(
workflow.task_group_relations,
{
Relation(from_task_id="task1_error", to_task_id="fail1", is_error=True),
Relation(from_task_id="task1_ok", to_task_id="task2", is_error=False),
Relation(from_task_id="task1_ok", to_task_id="task3", is_error=False),
Relation(from_task_id="task2_error", to_task_id="fail1", is_error=True),
Relation(from_task_id="task2_ok", to_task_id="task3", is_error=False),
Relation(from_task_id="task2_ok", to_task_id="task4_first", is_error=False),
Relation(from_task_id="task3_error", to_task_id="fail1", is_error=True),
Relation(from_task_id="task3_ok", to_task_id="end1", is_error=False),
Relation(from_task_id="task4_last", to_task_id="fail1", is_error=True),
Relation(from_task_id="task4_last", to_task_id="task1", is_error=False),
Relation(from_task_id="task4_last", to_task_id="task2", is_error=False),
Relation(from_task_id="task4_last", to_task_id="task3", is_error=False),
},
)
def test_last_task_id_of_error_flow(self):
task_group = TaskGroup(
name="TASK_GROUP",
error_downstream_name="AAAA",
tasks=[Task(task_id="TASK", template_name="dummy.tpl")],
)
task_group.add_state_handler_if_needed()
self.assertEqual("TASK_GROUP_error",
task_group.last_task_id_of_error_flow)
class WorkflowTemplateTestCase(TestCase, TemplateTestMixin):
TEMPLATE_NAME = "workflow.tpl"
DEFAULT_TEMPLATE_PARAMS = dict(
dag_name="test_dag",
dependencies={"import awesome_stuff"},
task_groups=[
TaskGroup(
name="TASK_GROUP",
tasks=[
Task(task_id="first_task", template_name="dummy.tpl"),
Task(task_id="second_task", template_name="dummy.tpl"),
],
)
],
job_properties={"user.name": "USER"},
config={},
relations={Relation(from_task_id="TASK_1", to_task_id="TASK_2")},
schedule_interval=3,
start_days_ago=3,
task_map={"oozie-task": ["airflow-task"]},
)
def test_green_path(self):
res = render_template(self.TEMPLATE_NAME, **self.DEFAULT_TEMPLATE_PARAMS)
self.assertValidPython(res)
def test_should_add_end_success_workflow_node(self):
# Given
transformer = AddWorkflowNotificationTransformer()
workflow = Workflow(input_directory_path="",
output_directory_path="",
dag_name="DAG_NAME_B")
props = PropertySet(
job_properties={
PROP_WORKFLOW_NOTIFICATION_URL: "http://example.com/workflow"
})
first_task_group = TaskGroup(
name="first_task",
tasks=[Task(task_id="first_task", template_name="dummy.tpl")])
# When
workflow.task_groups[first_task_group.name] = first_task_group
# Then
transformer.process_workflow_after_convert_nodes(workflow, props)
self.assertIn(END_SUCCESS_TASK_GROUP_NAME, workflow.task_groups.keys())
self.assertIn(END_SUCCESS_TASK_GROUP_NAME,
first_task_group.downstream_names)
self.assertEqual(
[
Task(
task_id=END_SUCCESS_TASK_GROUP_NAME,
template_name="http.tpl",
trigger_rule="one_success",
template_params={"url": "http://example.com/workflow"},
)
],
workflow.task_groups[END_SUCCESS_TASK_GROUP_NAME].tasks,
)
def test_last_task_id_of_ok_flow(self):
task_group = TaskGroup(
name="task1",
error_downstream_name="AAAA",
tasks=[Task(task_id="TASK", template_name="dummy.tpl")],
)
self.assertEqual("TASK", task_group.last_task_id_of_ok_flow)
class SubWorkflowTemplateTestCase(BaseTestCases.BaseTemplateTestCase):
TEMPLATE_NAME = "subworkflow.tpl"
DEFAULT_TEMPLATE_PARAMS = dict(
dependencies={"import awesome_stuff"},
task_groups=[
TaskGroup(
name="AAA",
tasks=[
Task(task_id="first_task", template_name="dummy.tpl"),
Task(task_id="second_task", template_name="dummy.tpl"),
],
relations=[
Relation(from_task_id="first_task",
to_task_id="second_task")
],
)
],
job_properties={"user.name": "USER"},
config={"key": "value"},
relations={Relation(from_task_id="TASK_1", to_task_id="TASK_2")},
)
def test_green_path(self):
res = render_template(self.TEMPLATE_NAME,
**self.DEFAULT_TEMPLATE_PARAMS)
self.assertValidPython(res)
def _create_notification_task_group(url_template, start_task_name, status):
url = (
url_template.replace("$parentId", "").replace("$jobId", "{{ run_id }}").replace("$status", status)
)
start_workflow_node = TaskGroup(
name=start_task_name,
tasks=[Task(task_id=start_task_name, template_name="http.tpl", template_params=dict(url=url))],
)
return start_workflow_node
def _create_workflow():
return Workflow(
dag_name="DAG_NAME",
input_directory_path="/tmp/input",
output_directory_path="/tmp/output",
task_group_relations={
Relation(from_task_id="DAG_NAME_A", to_task_id="DAG_NAME_B")
},
task_groups=dict(TASK_NAME=TaskGroup(
name="DAG_NAME_A",
tasks=[Task(task_id="task_name", template_name="dummy.tpl")])),
dependencies={"import IMPORT"},
)
def test_should_do_nothing_when_notification_url_not_configured(self):
# Given
transformer = AddWorkflowNotificationTransformer()
workflow = Workflow(input_directory_path="",
output_directory_path="",
dag_name="DAG_NAME_B")
props = PropertySet()
first_task_group = TaskGroup(
name="first_task",
tasks=[Task(task_id="first_task", template_name="dummy.tpl")])
workflow.task_groups[first_task_group.name] = first_task_group
# When
transformer.process_workflow_after_convert_nodes(workflow, props)
# Then
self.assertEqual({first_task_group.name}, workflow.task_groups.keys())
def convert_nodes(self):
"""
For each Oozie node, converts it into relations and internal relations.
It uses the mapper, which is stored in ParsedActionNode. The result is saved in ParsedActionNode.tasks
and ParsedActionNode.relations
"""
logging.info("Converting nodes to tasks and inner relations")
for name, p_node in self.workflow.nodes.copy().items():
tasks, relations = p_node.mapper.to_tasks_and_relations()
dependencies = p_node.mapper.required_imports()
p_node.tasks = tasks
p_node.relations = relations
self.workflow.task_groups[name] = TaskGroup(
name=name,
tasks=tasks,
relations=relations,
dependencies=dependencies,
downstream_names=p_node.downstream_names,
error_downstream_name=p_node.error_downstream_name,
)
del self.workflow.nodes[name]
def test_first_task_id(self):
task_group = TaskGroup(
name="TASK_GROUP",
tasks=[Task(task_id="TASK", template_name="dummy.tpl")])
self.assertEqual("TASK", task_group.first_task_id)
def setUp(self):
self.task_group = TaskGroup(
name="task1",
tasks=[Task(task_id="first_task_id", template_name="dummy.tpl")])