def test_should_keep_connected_nodes_in_error_state(self):
"""
Graph before:
.. graphviz::
digraph foo {
S -> A
}
Graph after:
.. graphviz::
digraph foo {
S -> A
}
Where:
A - first_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"
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)])
start_node = ParsedActionNode(
mapper=start_mapper,
tasks=[self._get_dummy_task(start_mapper.name)])
start_node.error_xml = first_node.name
workflow.nodes[first_mapper.name] = first_node
workflow.nodes[start_mapper.name] = start_node
transformer.process_workflow(workflow)
self.assertEqual({start_mapper.name, first_mapper.name},
set(workflow.nodes.keys()))
self.assertEqual([], start_node.downstream_names)
self.assertEqual([], first_node.downstream_names)
self.assertEqual(first_mapper.name, start_node.error_xml)
self.assertEqual(None, first_node.error_xml)
def test_should_remove_node_in_error_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"
second_mapper = mock.Mock(spec=KillMapper)
second_mapper.name = "second_task"
first_node = ParsedActionNode(first_mapper)
second_node = ParsedActionNode(second_mapper)
first_node.error_xml = second_mapper.name
workflow.nodes[first_mapper.name] = first_node
workflow.nodes[second_mapper.name] = second_node
transformer.process_workflow(workflow)
self.assertEqual({first_mapper.name}, set(workflow.nodes.keys()))