File size: 1,740 Bytes
88e3f4a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 | import pytest
from omniff.graph.executor import GraphExecutor
from omniff.graph.types import OmniGraph, OmniNode
from omniff.nodes.registry import NodeRegistry
def make_test_graph():
g = OmniGraph(id="test")
g.add_node(OmniNode(id="step1", node_type="echo", config={"value": "hello"}))
g.add_node(OmniNode(id="step2", node_type="echo", config={"value": "world"}))
g.add_node(OmniNode(id="step3", node_type="concat", config={}))
g.add_edge("step1", "step3")
g.add_edge("step2", "step3")
return g
def echo_node(inputs: dict, config: dict) -> dict:
return {"text": config["value"]}
def concat_node(inputs: dict, config: dict) -> dict:
texts = [v["text"] for v in inputs.values() if "text" in v]
return {"text": " ".join(texts)}
def test_executor_topological_order():
graph = make_test_graph()
registry = NodeRegistry()
registry.register("echo", echo_node)
registry.register("concat", concat_node)
executor = GraphExecutor(registry)
result = executor.execute(graph, {})
assert result["step3"]["text"] in ("hello world", "world hello")
def test_executor_single_node():
g = OmniGraph(id="single")
g.add_node(OmniNode(id="only", node_type="echo", config={"value": "test"}))
registry = NodeRegistry()
registry.register("echo", echo_node)
executor = GraphExecutor(registry)
result = executor.execute(g, {})
assert result["only"]["text"] == "test"
def test_executor_missing_node_type():
g = OmniGraph(id="bad")
g.add_node(OmniNode(id="n1", node_type="nonexistent", config={}))
registry = NodeRegistry()
executor = GraphExecutor(registry)
with pytest.raises(KeyError, match="nonexistent"):
executor.execute(g, {})
|