app/backend/tests/flow/test_runs.py

"""What a run is made of: isolation, the per-node record, and what it reports.

The service itself needs a database, so what is checked here is the part that
decides whether a run is correct — that two runs of one flow cannot see each
other's messages, that every node executed is reported once, and that the
parameters a caller sends are refused before anything runs if they are wrong.
"""

import pytest

from app.flow.messages import DType, MessageSpec
from app.flow.nodes import Node
from app.flow.pipeline import Pipeline
from app.flow.runs import (
    RunRejected,
    batch_issues,
    collect_result,
    digest_of,
    required_labels,
    seed_values,
)
from app.flow.schemas import FlowDef, FlowInput, NodeDef
from app.flow.state import MemoryState


def spec(name: str, dtype: DType = DType.FLOAT, **kwargs) -> MessageSpec:
    return MessageSpec(name=name, dtype=dtype, **kwargs)


def make_node(node_id: str, flow: str, f, requires=(), provides=()) -> Node:
    node = Node(f=f, requires=list(requires), provides=list(provides), name=node_id)
    node.assign_flow(flow, node_id)
    return node


def double_flow() -> FlowDef:
    """A flow with one input, one node and one declared output."""
    return FlowDef(
        name="study",
        mode="batch",
        inputs=[FlowInput(spec=spec("lr"), initial=0.1)],
        outputs=["loss"],
        nodes=[
            NodeDef(
                id="train",
                requires=[spec("lr")],
                provides=[spec("loss")],
            )
        ],
    )


def build(flow: FlowDef, state: MemoryState, observer=None) -> Pipeline:
    """The pipeline a run drives, without the controller that normally builds it."""
    node = make_node(
        "train",
        flow.name,
        lambda lr, params: {"loss": lr * 2},
        requires=[spec("lr")],
        provides=[spec("loss")],
    )
    return Pipeline(nodes=[node], state=state, observer=observer)


# -----------------------------------------------------------------------------
# Isolation — the reason a run has a state backend of its own
# -----------------------------------------------------------------------------


def test_two_runs_of_one_flow_do_not_see_each_other():
    flow = double_flow()
    first, second = MemoryState(), MemoryState()

    build(flow, first).run(seed_values(flow, {"lr": 0.5}))
    build(flow, second).run(seed_values(flow, {"lr": 4.0}))

    assert collect_result(flow, first) == {"loss": 1.0}
    assert collect_result(flow, second) == {"loss": 8.0}


def test_result_falls_back_to_everything_the_flow_holds():
    flow = double_flow()
    flow.outputs = []
    state = MemoryState()

    build(flow, state).run(seed_values(flow, {"lr": 1.0}))

    # The input is part of what the flow ended up holding; the engine's own
    # bookkeeping keys are not.
    assert collect_result(flow, state) == {"lr": 1.0, "loss": 2.0}


# -----------------------------------------------------------------------------
# The record — every node a run executed, reported once
# -----------------------------------------------------------------------------


def test_observer_sees_every_node_once():
    flow = double_flow()
    seen = []

    build(flow, MemoryState(), observer=seen.append).run(seed_values(flow, {"lr": 1.0}))

    assert [(o.node, o.ok) for o in seen] == [("study.train", True)]
    assert seen[0].outputs == 1


def test_observer_reports_a_failing_node_with_its_error():
    seen = []

    def boom(params):
        raise ValueError("no convergence")

    node = make_node("train", "study", boom, provides=[spec("loss")])
    Pipeline(nodes=[node], state=MemoryState(), observer=seen.append).run()

    assert len(seen) == 1
    assert not seen[0].ok
    assert "no convergence" in seen[0].error


def test_a_failing_observer_does_not_take_the_node_down():
    def refuse(_outcome):
        raise RuntimeError("the database is gone")

    node = make_node(
        "train", "study", lambda params: {"loss": 1.0}, provides=[spec("loss")]
    )
    pipeline = Pipeline(nodes=[node], state=MemoryState(), observer=refuse)
    pipeline.run()

    assert pipeline.state["study.loss"] == 1.0


# -----------------------------------------------------------------------------
# What a caller may ask for
# -----------------------------------------------------------------------------


def test_parameters_must_be_declared_inputs():
    flow = double_flow()
    with pytest.raises(RunRejected, match="not an input"):
        seed_values(flow, {"learning_rate": 0.1})


def test_parameters_are_type_checked_before_anything_runs():
    flow = double_flow()
    with pytest.raises(RunRejected, match="lr"):
        seed_values(flow, {"lr": "fast"})


def test_a_rate_limited_port_cannot_be_run_as_a_batch():
    flow = double_flow()
    flow.nodes[0].provides = [spec("loss", interval=30)]
    # Without a queue there is no timer to release what an interval holds, so
    # the value would be dropped rather than delayed.
    assert batch_issues(flow)
    assert not batch_issues(double_flow())


def test_the_digest_identifies_the_inputs_not_their_order():
    assert digest_of({"a": 1, "b": 2}, 3) == digest_of({"b": 2, "a": 1}, 3)
    assert digest_of({"a": 1}, 3) != digest_of({"a": 1}, 4)


def test_labels_come_from_the_nodes_that_ask_for_a_device():
    flow = double_flow()
    assert required_labels(flow) == []
    flow.nodes[0].device = "gpu"
    assert required_labels(flow) == ["gpu"]