import pytest
from typing import Type

from vellum.workflows.edges.edge import Edge
from vellum.workflows.graph.graph import Graph, NoPortsNode
from vellum.workflows.nodes.bases.base import BaseNode
from vellum.workflows.nodes.displayable.final_output_node import FinalOutputNode
from vellum.workflows.ports.port import Port
from vellum.workflows.state.base import BaseState
from vellum.workflows.triggers import ManualTrigger
from vellum.workflows.triggers.schedule import ScheduleTrigger


def test_graph__empty():
    # WHEN we create an empty graph
    graph = Graph.empty()

    # THEN the graph has no entrypoints
    assert len(list(graph.entrypoints)) == 0

    # AND no nodes
    assert len(list(graph.nodes)) == 0

    # AND no edges
    assert len(list(graph.edges)) == 0

    # AND string representation indicates empty graph
    assert str(graph) == "Graph(empty)"


def test_graph__from_node():
    # GIVEN a node
    class MyNode(BaseNode):
        pass

    # WHEN we create a graph from the node
    graph = Graph.from_node(MyNode)

    # THEN the graph has the node as the entrypoint
    assert set(graph.entrypoints) == {MyNode}

    # AND one node
    assert len(list(graph.nodes)) == 1

    # AND no edges
    assert len(list(graph.edges)) == 0


def test_graph__from_edge():
    # GIVEN an edge
    class SourceNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    edge = Edge(SourceNode.Ports.default, TargetNode)

    # WHEN we create a graph from the edge
    graph = Graph.from_edge(edge)

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND two nodes
    assert len(list(graph.nodes)) == 2

    # AND one edge
    assert len(list(graph.edges)) == 1


def test_graph__node_to_node():
    # GIVEN two nodes
    class SourceNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph from the source node to the target node
    graph = SourceNode >> TargetNode

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND two nodes
    assert len(list(graph.nodes)) == 2

    # AND one edge
    assert len(list(graph.edges)) == 1


def test_graph__port_to_node():
    # GIVEN a two nodes, where the source node has a port
    class SourceNode(BaseNode):
        class Ports(BaseNode.Ports):
            output = Port.on_else()

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph from the source port to the target node
    graph = SourceNode.Ports.output >> TargetNode

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND two nodes
    assert len(list(graph.nodes)) == 2

    # AND one edge
    assert len(list(graph.edges)) == 1


def test_graph__graph_to_node():
    # GIVEN three nodes
    class SourceNode(BaseNode):
        pass

    class MiddleNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph from the source node to the target node via the middle node
    graph = SourceNode >> MiddleNode >> TargetNode

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 2


def test_graph__edgeless_graph_to_node():
    # GIVEN two nodes
    class SourceNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # AND a graph of the first node
    subgraph = Graph.from_node(SourceNode)

    # WHEN we create a graph from the source graph to the target node
    graph = subgraph >> TargetNode

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND two nodes
    assert len(list(graph.nodes)) == 2

    # AND one edge
    assert len(list(graph.edges)) == 1

    # AND the first node's port has reference to the edge
    assert set(SourceNode.Ports.default.edges) == set(graph.edges)


def test_graph__edgeless_graph_to_graph():
    # GIVEN two nodes
    class SourceNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # AND a graph of the first node
    subgraph = Graph.from_node(SourceNode)

    # AND a graph of the second node
    target_subgraph = Graph.from_node(TargetNode)

    # WHEN we create a graph from the source graph to the target node
    graph = subgraph >> target_subgraph

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND two nodes
    assert len(list(graph.nodes)) == 2

    # AND one edge
    assert len(list(graph.edges)) == 1

    # AND the first node's port has reference to the edges
    assert set(SourceNode.Ports.default.edges) == set(graph.edges)


def test_graph__graph_to_edgeless_graph():
    # GIVEN three nodes
    class SourceNode(BaseNode):
        pass

    class MiddleNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # AND a graph of the first two nodes
    subgraph = SourceNode >> MiddleNode

    # AND a graph of the third node
    target_subgraph = Graph.from_node(TargetNode)

    # WHEN we create a graph from the source graph to the target node
    graph = subgraph >> target_subgraph

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 2


def test_graph__node_to_graph():
    # GIVEN three nodes
    class SourceNode(BaseNode):
        pass

    class MiddleNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # AND a graph of the last two nodes
    target_subgraph = MiddleNode >> TargetNode

    # WHEN we create a graph from the source node to the target graph
    graph = SourceNode >> target_subgraph

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 2


def test_graph__repeated_edge():
    # GIVEN two nodes
    class SourceNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph from the source node to the target node twice
    graph = SourceNode >> TargetNode >> SourceNode >> TargetNode

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND two nodes
    assert len(list(graph.nodes)) == 2

    # AND two edges
    assert len(list(graph.edges)) == 2

    # AND the first node's port has reference to just one edge
    assert len(list(SourceNode.Ports.default.edges)) == 1


def test_graph__node_to_set():
    # GIVEN three nodes, one with ports
    class SourceNode(BaseNode):
        pass

    class MiddleNode(BaseNode):
        class Ports(BaseNode.Ports):
            top = Port.on_if(SourceNode.Execution.count.less_than(1))
            bottom = Port.on_else()

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph with a set
    graph = SourceNode >> {
        MiddleNode.Ports.top >> SourceNode,
        MiddleNode.Ports.bottom >> TargetNode,
    }

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 3


def test_graph__graph_to_set():
    # GIVEN four nodes, one with ports
    class SourceNode(BaseNode):
        pass

    class SecondNode(BaseNode):
        pass

    class MiddleNode(BaseNode):
        class Ports(BaseNode.Ports):
            top = Port.on_if(SourceNode.Execution.count.less_than(1))
            bottom = Port.on_else()

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph between a graph and a set
    graph = (
        SourceNode
        >> SecondNode
        >> {
            MiddleNode.Ports.top >> SourceNode,
            MiddleNode.Ports.bottom >> TargetNode,
        }
    )

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 4

    # AND two edges
    assert len(list(graph.edges)) == 4


def test_graph__graph_set_to_set():
    # GIVEN five nodes
    class SourceNode(BaseNode):
        pass

    class TopNode(BaseNode):
        pass

    class BottomNode(BaseNode):
        pass

    class ConditionalNode(BaseNode):
        class Ports(BaseNode.Ports):
            top = Port.on_if(SourceNode.Execution.count.less_than(1))
            bottom = Port.on_else()

    class EndNode(BaseNode):
        pass

    # WHEN we create a graph that draws an edge between sets
    graph = (
        SourceNode
        >> {
            TopNode,
            BottomNode,
        }
        >> {
            ConditionalNode.Ports.top >> SourceNode,
            ConditionalNode.Ports.bottom >> EndNode,
        }
    )

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 5

    # AND two edges
    assert len(list(graph.edges)) == 6


def test_graph__port_to_graph():
    # GIVEN three nodes where the first node has a port
    class SourceNode(BaseNode):
        class Ports(BaseNode.Ports):
            custom = Port.on_else()

    class MiddleNode(BaseNode):
        pass

    class EndNode(BaseNode):
        pass

    # WHEN we create a graph from the port to a subgraph
    graph = SourceNode.Ports.custom >> (MiddleNode >> EndNode)

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 2


def test_graph__port_to_set():
    # GIVEN three nodes where the first node has a port
    class SourceNode(BaseNode):
        class Ports(BaseNode.Ports):
            custom = Port.on_else()

    class TopNode(BaseNode):
        pass

    class BottomNode(BaseNode):
        pass

    # WHEN we create a graph from the port to the set
    graph = SourceNode.Ports.custom >> {
        TopNode,
        BottomNode,
    }

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 2


def test_graph__set_to_node():
    # GIVEN three nodes
    class TopNode(BaseNode):
        pass

    class BottomNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph from a set to a node
    graph = {
        TopNode,
        BottomNode,
    } >> TargetNode

    # THEN the graph has both the top node and the bottom node as the entrypoints
    assert set(graph.entrypoints) == {TopNode, BottomNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 2


def test_graph__node_to_port():
    # GIVEN two nodes, one with a port
    class SourceNode(BaseNode):
        pass

    class MiddleNode(BaseNode):
        class Ports(BaseNode.Ports):
            custom = Port.on_else()

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph from the source node to the target node
    graph = SourceNode >> MiddleNode.Ports.custom >> TargetNode

    # THEN the graph has the source node as the entrypoint
    assert set(graph.entrypoints) == {SourceNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 2


def test_graph__set_to_graph():
    # GIVEN three nodes
    class SourceNode(BaseNode):
        pass

    class MiddleNode(BaseNode):
        pass

    class TargetNode(BaseNode):
        pass

    # WHEN we create a graph from a set to a graph
    graph: Graph = {SourceNode, MiddleNode} >> Graph.from_node(TargetNode)

    # THEN the graph has the source node and middle node as the entrypoints
    assert set(graph.entrypoints) == {SourceNode, MiddleNode}

    # AND three nodes
    assert len(list(graph.nodes)) == 3

    # AND two edges
    assert len(list(graph.edges)) == 2


def test_graph__str_simple_linear():
    # GIVEN a simple linear graph: A -> B -> C
    class NodeA(BaseNode):
        pass

    class NodeB(BaseNode):
        pass

    class NodeC(BaseNode):
        pass

    graph = NodeA >> NodeB >> NodeC

    # WHEN we convert the graph to string
    result = str(graph)

    # THEN it shows the linear flow structure
    expected_lines = ["Graph:", "  └─ NodeA", "     └─ NodeB", "        └─ NodeC"]
    assert result == "\n".join(expected_lines)


def test_graph__str_with_branching():
    # GIVEN a graph with branching: A -> {B, C}
    class NodeA(BaseNode):
        pass

    class NodeB(BaseNode):
        pass

    class NodeC(BaseNode):
        pass

    graph = NodeA >> {NodeB, NodeC}

    # WHEN we convert the graph to string
    result = str(graph)

    # THEN it shows the branching structure
    # Note: The order might vary due to set ordering, so we check for the structure
    lines = result.split("\n")
    assert lines[0] == "Graph:"
    assert lines[1] == "  └─ NodeA"

    # Should have two branches (order may vary)
    branch_lines = [line.strip() for line in lines[2:] if line.strip()]
    assert len(branch_lines) == 2
    assert any("NodeB" in line for line in branch_lines)
    assert any("NodeC" in line for line in branch_lines)
    assert all(line.startswith("├─ ") or line.startswith("└─ ") for line in branch_lines)


def test_graph__str_with_loop():
    # GIVEN a graph with a loop: A -> B -> A (loop back)
    class NodeA(BaseNode):
        pass

    class NodeB(BaseNode):
        pass

    # Create the loop manually using edges
    edge1 = Edge(NodeA.Ports.default, NodeB)
    edge2 = Edge(NodeB.Ports.default, NodeA)

    graph = Graph(entrypoints={NodeA.Ports.default}, edges=[edge1, edge2], terminals={NodeA.Ports.default})

    # WHEN we convert the graph to string
    result = str(graph)

    # THEN it shows the loop with cycle detection
    expected_lines = ["Graph:", "  └─ NodeA", "     └─ NodeB", "        └─ NodeA ⟲ (loops back)"]
    assert result == "\n".join(expected_lines)


def test_graph__str_empty_graph():
    # GIVEN an empty graph
    graph = Graph(entrypoints=set(), edges=[], terminals=set())

    # WHEN we convert the graph to string
    result = str(graph)

    # THEN it shows empty graph message
    assert result == "Graph(empty)"


def test_graph__str_single_node():
    # GIVEN a graph with just one node and no edges
    class SingleNode(BaseNode):
        pass

    graph = Graph.from_node(SingleNode)

    # WHEN we convert the graph to string
    result = str(graph)

    # THEN it shows the single node
    assert "SingleNode.default" in result
    assert "Graph:" in result


def test_graph__from_node_with_empty_ports():
    """
    Tests that building a graph from a single node with empty Ports class generates 1 node.
    """

    # GIVEN a node with an empty Ports class
    class NodeWithEmptyPorts(BaseNode):
        class Ports(BaseNode.Ports):
            pass

    # WHEN we create a graph from the node
    graph = Graph.from_node(NodeWithEmptyPorts)

    # THEN the graph should have exactly 1 node
    assert len(list(graph.nodes)) == 1


def test_graph__manual_trigger_to_node():
    # GIVEN a node
    class MyNode(BaseNode):
        pass

    # WHEN we create graph with ManualTrigger >> Node (class-level, no instantiation)
    graph = ManualTrigger >> MyNode

    # THEN the graph has one trigger edge
    trigger_edges = list(graph.trigger_edges)
    assert len(trigger_edges) == 1
    assert trigger_edges[0].trigger_class == ManualTrigger
    assert trigger_edges[0].to_node == MyNode

    # AND the graph has one trigger
    triggers = list(graph.triggers)
    assert len(triggers) == 1
    assert triggers[0] == ManualTrigger

    # AND the graph has one node
    assert len(list(graph.nodes)) == 1
    assert MyNode in list(graph.nodes)


def test_graph__manual_trigger_to_set_of_nodes():
    # GIVEN two nodes
    class NodeA(BaseNode):
        pass

    class NodeB(BaseNode):
        pass

    # WHEN we create graph with ManualTrigger >> {NodeA, NodeB}
    graph = ManualTrigger >> {NodeA, NodeB}

    # THEN the graph has two trigger edges
    trigger_edges = list(graph.trigger_edges)
    assert len(trigger_edges) == 2

    # AND both edges connect to the same ManualTrigger class
    assert all(edge.trigger_class == ManualTrigger for edge in trigger_edges)

    # AND edges connect to both nodes
    target_nodes = {edge.to_node for edge in trigger_edges}
    assert target_nodes == {NodeA, NodeB}

    # AND the graph has one unique trigger
    triggers = list(graph.triggers)
    assert len(triggers) == 1

    # AND the graph has two nodes
    assert len(list(graph.nodes)) == 2


def test_graph__manual_trigger_to_graph():
    # GIVEN a graph of nodes
    class NodeA(BaseNode):
        pass

    class NodeB(BaseNode):
        pass

    node_graph = NodeA >> NodeB

    # WHEN we create graph with ManualTrigger >> Graph
    graph = ManualTrigger >> node_graph

    # THEN the graph has a trigger edge to the entrypoint
    trigger_edges = list(graph.trigger_edges)
    assert len(trigger_edges) == 1
    assert trigger_edges[0].to_node == NodeA

    # AND the graph preserves the original edges
    edges = list(graph.edges)
    assert len(edges) == 1
    assert edges[0].to_node == NodeB

    # AND the graph has both nodes
    nodes = list(graph.nodes)
    assert len(nodes) == 2
    assert NodeA in nodes
    assert NodeB in nodes


def test_graph__manual_trigger_to_set_of_graphs_preserves_edges():
    # GIVEN two graphs of nodes
    class NodeA(BaseNode):
        pass

    class NodeB(BaseNode):
        pass

    class NodeC(BaseNode):
        pass

    class NodeD(BaseNode):
        pass

    graph_one = NodeA >> NodeB
    graph_two = NodeC >> NodeD

    # WHEN we create a graph with ManualTrigger >> {Graph1, Graph2}
    combined_graph = ManualTrigger >> {graph_one, graph_two}

    # THEN the combined graph has trigger edges to both entrypoints
    trigger_edges = list(combined_graph.trigger_edges)
    assert len(trigger_edges) == 2
    assert {edge.to_node for edge in trigger_edges} == {NodeA, NodeC}

    # AND the combined graph preserves all downstream edges
    edges = list(combined_graph.edges)
    assert len(edges) == 2
    assert {(edge.from_port.node_class, edge.to_node) for edge in edges} == {
        (NodeA, NodeB),
        (NodeC, NodeD),
    }

    # AND the combined graph still exposes all nodes
    nodes = list(combined_graph.nodes)
    assert {NodeA, NodeB, NodeC, NodeD}.issubset(nodes)


def test_graph__node_to_trigger_raises():
    # GIVEN a node and trigger
    class MyNode(BaseNode):
        pass

    # WHEN we try to create Node >> Trigger (class-level)
    # THEN it raises TypeError
    with pytest.raises(TypeError, match="Cannot create edge targeting trigger"):
        MyNode >> ManualTrigger

    # WHEN we try to create Node >> Trigger (instance-level)
    # THEN it also raises TypeError
    with pytest.raises(TypeError, match="Cannot create edge targeting trigger"):
        MyNode >> ManualTrigger


def test_graph__trigger_then_graph_then_node():
    # GIVEN a trigger, a node, and another node
    class StartNode(BaseNode):
        pass

    class EndNode(BaseNode):
        pass

    # WHEN we create Trigger >> Node >> Node
    graph = ManualTrigger >> StartNode >> EndNode

    # THEN the graph has one trigger edge
    trigger_edges = list(graph.trigger_edges)
    assert len(trigger_edges) == 1
    assert trigger_edges[0].to_node == StartNode

    # AND the graph has one regular edge
    edges = list(graph.edges)
    assert len(edges) == 1
    assert edges[0].to_node == EndNode

    # AND the graph has both nodes
    nodes = list(graph.nodes)
    assert len(nodes) == 2


def test_graph__set_of_trigger_graphs_preserves_trigger_edges():
    """Test that combining graphs with triggers via a set preserves trigger edges.

    This tests the fix for Graph.from_set() not propagating _trigger_edges.
    """

    # GIVEN a custom scheduled trigger
    class MyScheduledTrigger(ScheduleTrigger):
        class Config:
            cron = "0 9 * * *"

    # AND two trigger-initiated graphs
    class MyFirstNode(BaseNode):
        pass

    class MySecondNode(BaseNode):
        pass

    trigger_graph_a = MyScheduledTrigger >> MyFirstNode
    trigger_graph_b = MyScheduledTrigger >> MySecondNode

    # WHEN we combine them in a set
    combined_graph = Graph.from_set({trigger_graph_a, trigger_graph_b})

    # THEN the combined graph has both trigger edges
    trigger_edges = list(combined_graph.trigger_edges)
    assert len(trigger_edges) == 2

    # AND both edges point to the correct nodes
    target_nodes = {edge.to_node for edge in trigger_edges}
    assert target_nodes == {MyFirstNode, MySecondNode}

    # AND the graph exposes the trigger
    triggers = list(combined_graph.triggers)
    assert len(triggers) == 1
    assert triggers[0] == MyScheduledTrigger


def test_graph__set_of_trigger_graphs_to_node_preserves_trigger_edges():
    # GIVEN a custom scheduled trigger
    class MyScheduledTrigger(ScheduleTrigger):
        class Config:
            cron = "0 9 * * *"

    # AND two trigger-initiated graphs
    class MyFirstNode(BaseNode):
        pass

    class MySecondNode(BaseNode):
        pass

    class MyThirdNode(BaseNode):
        pass

    trigger_graph_a = MyScheduledTrigger >> MyFirstNode
    trigger_graph_b = MyScheduledTrigger >> MySecondNode

    # WHEN we combine them in a set and connect to another node
    combined_graph = {trigger_graph_a, trigger_graph_b} >> MyThirdNode

    # THEN the combined graph has both trigger edges
    trigger_edges = list(combined_graph.trigger_edges)
    assert len(trigger_edges) == 2

    # AND both trigger edges point to the correct initial nodes
    trigger_target_nodes = {edge.to_node for edge in trigger_edges}
    assert trigger_target_nodes == {MyFirstNode, MySecondNode}

    # AND the graph has regular edges to the third node
    regular_edges = list(combined_graph.edges)
    assert len(regular_edges) == 2
    assert all(edge.to_node == MyThirdNode for edge in regular_edges)

    # AND the graph has all three nodes
    nodes = list(combined_graph.nodes)
    assert len(nodes) == 3
    assert set(nodes) == {MyFirstNode, MySecondNode, MyThirdNode}


def test_graph__graph_rshift_graph_preserves_trigger_edges():
    """Test that Graph >> Graph preserves trigger edges from the right-hand graph."""

    # GIVEN a regular graph and a trigger-initiated graph
    class StartNode(BaseNode):
        pass

    class TriggerNode(BaseNode):
        pass

    class EndNode(BaseNode):
        pass

    regular_graph = StartNode >> TriggerNode
    trigger_graph = ManualTrigger >> EndNode

    # WHEN we combine them with >>
    combined = regular_graph >> trigger_graph

    # THEN the combined graph has the trigger edge
    trigger_edges = list(combined.trigger_edges)
    assert len(trigger_edges) == 1
    assert trigger_edges[0].to_node == EndNode

    # AND the graph exposes the trigger
    triggers = list(combined.triggers)
    assert len(triggers) == 1


def test_graph__graph_rshift_set_of_trigger_graphs_preserves_trigger_edges():
    """Test that Graph >> {TriggerGraph1, TriggerGraph2} preserves trigger edges."""

    # GIVEN a regular graph and two trigger-initiated graphs
    class StartNode(BaseNode):
        pass

    class TriggerNodeA(BaseNode):
        pass

    class TriggerNodeB(BaseNode):
        pass

    regular_graph = Graph.from_node(StartNode)
    trigger_graph_a = ManualTrigger >> TriggerNodeA
    trigger_graph_b = ManualTrigger >> TriggerNodeB

    # WHEN we combine them with >> and a set
    combined = regular_graph >> {trigger_graph_a, trigger_graph_b}

    # THEN the combined graph has both trigger edges
    trigger_edges = list(combined.trigger_edges)
    assert len(trigger_edges) == 2

    # AND both trigger edges point to the correct nodes
    trigger_target_nodes = {edge.to_node for edge in trigger_edges}
    assert trigger_target_nodes == {TriggerNodeA, TriggerNodeB}


class TestFinalOutputNoPortsNode:
    """Test that we preserve final output nodes as NoPortsNode when constructing a Graph."""

    def validate_graph(self, graph: Graph, final_output_node: Type[FinalOutputNode]):
        assert len(list(graph._terminals)) == 1
        node = list(graph._terminals)[0]
        assert isinstance(node, NoPortsNode)
        assert node.node_class == final_output_node

    def test_from_edge(self):
        # GIVEN
        class StartNode(BaseNode[BaseState]):
            pass

        class MyFinalOutput(FinalOutputNode[BaseState, str]):
            pass

        # WHEN
        edge = Edge(from_port=StartNode.Ports.default, to_node=MyFinalOutput)
        graph = Graph.from_edge(edge)

        # THEN
        self.validate_graph(graph, MyFinalOutput)

    def test_rshift__node(self):
        # GIVEN
        class StartNode(BaseNode[BaseState]):
            pass

        class MyFinalOutput(FinalOutputNode[BaseState, str]):
            pass

        # WHEN
        initial_graph = Graph.from_node(StartNode)
        graph = initial_graph >> MyFinalOutput

        # THEN
        self.validate_graph(graph, MyFinalOutput)

    def test_rshift__graph(self):
        # GIVEN
        class StartNode(BaseNode[BaseState]):
            pass

        class MyFinalOutput(FinalOutputNode[BaseState, str]):
            pass

        # WHEN
        subgraph = Graph.from_node(MyFinalOutput)
        graph = Graph.from_node(StartNode) >> subgraph

        # THEN
        self.validate_graph(subgraph, MyFinalOutput)
        self.validate_graph(graph, MyFinalOutput)

    def test_rshift__set_with_node(self):
        # GIVEN
        class StartNode(BaseNode[BaseState]):
            pass

        class MyFinalOutput(FinalOutputNode[BaseState, str]):
            pass

        # WHEN
        graph = Graph.from_node(StartNode) >> {MyFinalOutput}

        # THEN
        self.validate_graph(graph, MyFinalOutput)

    def test_rshift__set_with_graph(self):
        # GIVEN
        class StartNode(BaseNode[BaseState]):
            pass

        class MyFinalOutput(FinalOutputNode[BaseState, str]):
            pass

        # WHEN
        subgraph = Graph.from_node(MyFinalOutput)
        graph = Graph.from_node(StartNode) >> {subgraph}

        # THEN
        self.validate_graph(subgraph, MyFinalOutput)
        self.validate_graph(graph, MyFinalOutput)
