def test_should_keep_node_in_correct_flow(self):
transformer = RemoveKillTransformer()
workflow = Workflow(input_directory_path="",
output_directory_path="",
dag_name="DAG_NAME_B")
first_mapper = mock.Mock(spec=BaseMapper)
first_mapper.name = "first_task"
third_mapper = mock.Mock(spec=KillMapper)
third_mapper.name = "third_task"
first_node = ParsedActionNode(first_mapper)
third_node = ParsedActionNode(third_mapper)
first_node.downstream_names = [third_mapper.name]
workflow.nodes[first_mapper.name] = first_node
workflow.nodes[third_mapper.name] = third_node
transformer.process_workflow(workflow)
self.assertEqual({first_mapper.name, third_mapper.name},
set(workflow.nodes.keys()))
self.assertEqual([third_node.name], first_node.downstream_names)
self.assertEqual([], third_node.downstream_names)
def test_should_remove_end_node(self):
transformer = RemoveEndTransformer()
workflow = Workflow(input_directory_path="",
output_directory_path="",
dag_name="DAG_NAME_B")
first_mapper = mock.Mock(spec=BaseMapper)
first_mapper.name = "first_task"
end_mapper = mock.Mock(spec=EndMapper)
end_mapper.name = "second_task"
first_node = ParsedActionNode(mapper=first_mapper)
end_node = ParsedActionNode(mapper=end_mapper)
first_node.downstream_names = [end_node.name]
workflow.nodes[first_mapper.name] = first_node
workflow.nodes[end_mapper.name] = end_node
transformer.process_workflow(workflow)
self.assertEqual({first_mapper.name}, set(workflow.nodes.keys()))
self.assertEqual([], first_node.downstream_names)
def test_should_not_remove_end_node_when_connected_with_decision(self):
transformer = RemoveEndTransformer()
workflow = Workflow(input_directory_path="",
output_directory_path="",
dag_name="DAG_NAME_B")
decision_mapper = mock.Mock(spec=DecisionMapper)
decision_mapper.name = "first_task"
second_mapper = mock.Mock(spec=BaseMapper)
second_mapper.name = "second_task"
end_mapper = mock.Mock(spec=EndMapper)
end_mapper.name = "end_task"
decision_node = ParsedActionNode(
mapper=decision_mapper,
tasks=[self._get_dummy_task(decision_mapper.name)])
second_node = ParsedActionNode(
mapper=second_mapper,
tasks=[self._get_dummy_task(second_mapper.name)])
end_node = ParsedActionNode(
mapper=end_mapper, tasks=[self._get_dummy_task(end_mapper.name)])
decision_node.downstream_names = [second_mapper.name, end_mapper.name]
workflow.nodes[decision_mapper.name] = decision_node
workflow.nodes[second_mapper.name] = second_node
workflow.nodes[end_mapper.name] = end_node
transformer.process_workflow(workflow)
self.assertEqual(
{decision_mapper.name, second_mapper.name, end_mapper.name},
set(workflow.nodes.keys()))
self.assertEqual([second_mapper.name, end_mapper.name],
decision_node.downstream_names)
self.assertEqual([], second_node.downstream_names)
self.assertEqual([], end_node.downstream_names)
def test_should_keep_connected_nodes_in_correct_flow(self):
"""
Graph before:
.. graphviz::
digraph foo {
S -> A
A -> B
}
Graph after:
.. graphviz::
digraph foo {
S -> A
A -> B
}
Where:
A - first_task
B - second_task
S - start_task
"""
transformer = RemoveInaccessibleNodeTransformer()
workflow = Workflow(input_directory_path="",
output_directory_path="",
dag_name="DAG_NAME_B")
first_mapper = mock.Mock(spec=BaseMapper)
first_mapper.name = "first_task"
second_mapper = mock.Mock(spec=BaseMapper)
second_mapper.name = "second_task"
start_mapper = mock.Mock(spec=StartMapper)
start_mapper.name = "start_task"
first_node = ParsedActionNode(
mapper=first_mapper,
tasks=[self._get_dummy_task(first_mapper.name)])
second_node = ParsedActionNode(
mapper=second_mapper,
tasks=[self._get_dummy_task(second_mapper.name)])
start_node = ParsedActionNode(
mapper=start_mapper,
tasks=[self._get_dummy_task(start_mapper.name)])
start_node.downstream_names = [first_mapper.name]
first_node.downstream_names = [second_node.name]
workflow.nodes[first_mapper.name] = first_node
workflow.nodes[second_mapper.name] = second_node
workflow.nodes[start_mapper.name] = start_node
transformer.process_workflow(workflow)
self.assertEqual(
{start_mapper.name, first_mapper.name, second_mapper.name},
set(workflow.nodes.keys()))
self.assertEqual([first_mapper.name], start_node.downstream_names)
self.assertEqual([second_mapper.name], first_node.downstream_names)
self.assertEqual([], second_node.downstream_names)