from plesna.graph.graph import Graph
from plesna.graph.graph_set import GraphSet
from plesna.models.graphs import Edge, EdgeOnSet, Node


def test_init():
    graph_set = GraphSet()

    nodeA = Node(name="A")
    nodeB = Node(name="B")
    nodeC = Node(name="C")
    edge1 = EdgeOnSet(arrow="arrow", sources=[nodeA, nodeB], targets=[nodeC])

    graph_set.append(edge1)

    assert graph_set.node_sets == {frozenset([nodeA, nodeB]), frozenset([nodeC])}


def test_to_graph():
    graph_set = GraphSet()

    nodeA = Node(name="A")
    nodeB = Node(name="B")
    nodeC = Node(name="C")
    nodeD = Node(name="D")
    edge1 = EdgeOnSet(arrow="arrow-AB-C", sources=[nodeA, nodeB], targets=[nodeC])
    edge2 = EdgeOnSet(arrow="arrow-C-D", sources=[nodeC], targets=[nodeD])

    graph_set.append(edge1)
    graph_set.append(edge2)

    graph = graph_set.to_graph()
    assert graph.nodes == {
        nodeA,
        nodeB,
        nodeC,
        nodeD,
    }
    assert graph.edges == [
        Edge(arrow="arrow-AB-C", source=nodeA, target=nodeC),
        Edge(arrow="arrow-AB-C", source=nodeB, target=nodeC),
        Edge(arrow="arrow-C-D", source=nodeC, target=nodeD),
    ]