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