app/backend/tests/flow/test_replace_flow.py

"""Swapping one flow's nodes into a live graph, leaving the rest connected."""

import threading

from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes import Node
from fluksio.flow.pipeline import Pipeline
from tests.flow.test_pipeline import make_node, spec


def _shape(pipeline: Pipeline) -> tuple[dict[str, list[str]], dict[str, set[str]]]:
    """The wiring by id, so two pipelines built differently can be compared."""
    return (
        {msg: [n.id for n in nodes] for msg, nodes in pipeline.produces.items()},
        {
            node.id: {dep.id for dep in deps}
            for node, deps in pipeline.dependencies.items()
        },
    )


def test_replacing_a_flow_leaves_the_other_flows_nodes_alone():
    b_producer = make_node(
        "meter", "b", lambda params: {"power": 1.0}, provides=[spec("power")]
    )
    b_consumer = make_node(
        "load", "b", lambda power, params: None, requires=[spec("power")]
    )
    a_old = make_node("calc", "a", lambda params: None)
    pipeline = Pipeline(nodes=[a_old, b_producer, b_consumer])

    pipeline.replace_flow("a", [make_node("calc", "a", lambda params: None)])

    assert pipeline.get_node_by_id("b.meter") is b_producer
    assert pipeline.get_node_by_id("b.load") is b_consumer
    assert pipeline.dependencies[b_consumer] == frozenset({b_producer})


def test_a_consumer_in_another_flow_picks_up_the_new_producer():
    seen: list[float] = []
    consumer = make_node(
        "load",
        "b",
        lambda temp, params: seen.append(temp),
        requires=[spec("a.temp")],
    )
    pipeline = Pipeline(
        nodes=[
            make_node(
                "old", "a", lambda params: {"temp": 1.0}, provides=[spec("temp")]
            ),
            consumer,
        ]
    )

    replacement = make_node(
        "new", "a", lambda params: {"temp": 2.0}, provides=[spec("temp")]
    )
    pipeline.replace_flow("a", [replacement])

    assert pipeline.dependencies[consumer] == frozenset({replacement})
    replacement.inject()
    assert seen == [2.0]


def test_a_consumer_loses_its_edge_when_the_producer_flow_stops_providing():
    """The crux: the consumer becomes a root, not an orphan nothing can reach."""
    consumer = make_node(
        "load", "b", lambda temp, params: None, requires=[spec("a.temp")]
    )
    pipeline = Pipeline(
        nodes=[
            make_node(
                "meter", "a", lambda params: {"temp": 1.0}, provides=[spec("temp")]
            ),
            consumer,
        ]
    )

    pipeline.replace_flow("a", [make_node("quiet", "a", lambda params: None)])

    assert "a.temp" not in pipeline.produces
    assert pipeline.dependencies[consumer] == frozenset()
    assert consumer in pipeline._topological_sort()


def test_removing_a_flow_takes_its_nodes_out():
    ran: list[str] = []
    a_node = make_node(
        "meter",
        "a",
        lambda params: ran.append("a") or {"temp": 1.0},
        provides=[spec("temp")],
    )
    b_node = make_node(
        "beat",
        "b",
        lambda params: ran.append("b") or {"tick": 1.0},
        provides=[spec("tick")],
    )
    pipeline = Pipeline(nodes=[a_node, b_node])

    pipeline.remove_flow("a")

    assert pipeline.flow_nodes("a") == set()
    assert pipeline.get_node_by_id("a.meter") is None
    assert a_node not in pipeline.dependencies
    b_node.inject()
    assert ran == ["b"]

    # A caller still holding the removed node — a wave that picked its targets
    # before the removal — runs nothing rather than raising or reviving it.
    pipeline.run({}, nodes={a_node})
    assert ran == ["b"]


def test_a_replace_leaves_the_graph_a_rebuild_would_have_built():
    """A replace and a fresh build over the same nodes must agree, edge for edge."""

    def flow_a(tag: str) -> list[Node]:
        # Reads a message flow b owns, writes one flow b reads, and shares a
        # third with a producer in flow b — so the comparison covers edges in
        # both directions across the boundary, and the order of a fan-in whose
        # producers sit on either side of it.
        return [
            make_node(
                "meter",
                "a",
                lambda power, params: {"temp": 1.0},
                requires=[spec("b.power")],
                provides=[spec("temp"), MessageSpec(name=tag, dtype=DType.FLOAT)],
            ),
            make_node(
                "inner",
                "a",
                lambda temp, params: None,
                requires=[spec("temp")],
            ),
            make_node(
                "mirror",
                "a",
                lambda params: {"power": 1.0},
                provides=[MessageSpec(name="b.power", dtype=DType.FLOAT)],
            ),
        ]

    def flow_b() -> list[Node]:
        return [
            make_node(
                "supply", "b", lambda params: {"power": 1.0}, provides=[spec("power")]
            ),
            make_node(
                "load", "b", lambda temp, params: None, requires=[spec("a.temp")]
            ),
        ]

    # Both flows write it, so the order the two producers come out in is the
    # node list's, and a replace that appends rather than splices gets it wrong.
    b_nodes = flow_b()
    replaced = Pipeline(nodes=flow_a("old") + b_nodes)
    replaced.replace_flow("a", flow_a("new"))

    fresh = Pipeline(nodes=flow_a("new") + flow_b())

    assert _shape(replaced) == _shape(fresh)
    assert [n.id for n in replaced._topological_sort()] == [
        n.id for n in fresh._topological_sort()
    ]


def test_a_cascade_running_during_a_replace_finishes_on_the_graph_it_started_with():
    started = threading.Event()
    release = threading.Event()
    ran: list[str] = []

    def slow(params):
        started.set()
        release.wait(5)
        return {"temp": 1.0}

    source = make_node("source", "a", slow, provides=[spec("temp")])
    consumer = make_node(
        "load", "a", lambda temp, params: ran.append("load"), requires=[spec("temp")]
    )
    pipeline = Pipeline(nodes=[source, consumer])

    failure: list[BaseException] = []

    def wave() -> None:
        try:
            pipeline.run({})
        except BaseException as exc:  # noqa: BLE001 - reported, not swallowed
            failure.append(exc)

    thread = threading.Thread(target=wave)
    thread.start()
    # The replace lands while the wave is inside the executor holding the old
    # maps, which is the only moment the swap can be seen half-applied.
    assert started.wait(5)
    pipeline.replace_flow("a", [make_node("source", "a", lambda params: None)])
    release.set()
    thread.join(10)

    assert not thread.is_alive()
    assert failure == []
    assert ran == ["load"]


def test_replacing_a_flow_clears_its_pause_and_leaves_another_flows():
    pipeline = Pipeline(
        nodes=[
            make_node("one", "a", lambda params: None),
            make_node("two", "b", lambda params: None),
        ]
    )
    pipeline.pause("a")
    pipeline.pause("b")

    pipeline.replace_flow("a", [make_node("one", "a", lambda params: None)])

    assert pipeline.paused_flows() == ["b"]