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)