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, {})