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])
chain = Chain()
chain.add_node(final)
chain.delete_node(third)
assert len(chain.nodes) == 3
assert first in chain.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])
chain = Chain()
chain.add_node(final)
# when
chain.delete_node(first)
new_primary_node = [
node for node in chain.nodes if node.operation.operation_type == 'knn'
][0]
# then
assert len(chain.nodes) == 3
assert isinstance(new_primary_node, PrimaryNode)
