def test_delete_node_with_redirection():
first = PrimaryNode(operation_type='logit')
second = PrimaryNode(operation_type='lda')
third = SecondaryNode(operation_type='knn', nodes_from=[first, second])
final = SecondaryNode(operation_type='xgboost',
nodes_from=[third])
pipeline = Pipeline()
pipeline.add_node(final)
pipeline.delete_node(third)
assert len(pipeline.nodes) == 3
assert first in pipeline.root_node.nodes_from
def test_delete_primary_node():
# given
first = PrimaryNode(operation_type='logit')
second = PrimaryNode(operation_type='lda')
third = SecondaryNode(operation_type='knn', nodes_from=[first])
final = SecondaryNode(operation_type='xgboost',
nodes_from=[second, third])
pipeline = Pipeline(final)
# when
pipeline.delete_node(first)
new_primary_node = [node for node in pipeline.nodes if node.operation.operation_type == 'knn'][0]
# then
assert len(pipeline.nodes) == 3
assert isinstance(new_primary_node, PrimaryNode)
