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)