pyflowx/tests/test_loop.py

"""LoopSpec 任务循环展开测试。"""

from __future__ import annotations

import pyflowx as px
from pyflowx import Graph, LoopSpec, TaskSpec, task


def test_loop_range_expands_to_instances() -> None:
    """LoopSpec.range 展开为多实例，item 追加到 args。"""

    def process(item: int) -> int:
        return item * 2

    spec = TaskSpec("process", fn=process, loop=LoopSpec.range(3))
    graph = Graph.from_specs([spec])

    assert "process" not in graph.specs
    assert set(graph.specs.keys()) == {"process_0", "process_1", "process_2"}
    assert graph.specs["process_0"].args == (0,)
    assert graph.specs["process_2"].args == (2,)
    assert graph.specs["process_0"].loop is None


def test_loop_from_iterable() -> None:
    """LoopSpec.from_iterable 接受任意可迭代对象。"""

    def echo(item: str) -> str:
        return item

    spec = TaskSpec("echo", fn=echo, loop=LoopSpec.from_iterable(["a", "b"]))
    graph = Graph.from_specs([spec])

    assert set(graph.specs.keys()) == {"echo_0", "echo_1"}
    assert graph.specs["echo_0"].args == ("a",)
    assert graph.specs["echo_1"].args == ("b",)


def test_loop_range_with_start_step() -> None:
    """LoopSpec.range 支持 start/stop/step 三参数（与 Python range 一致）。"""

    def process(item: int) -> int:
        return item

    spec = TaskSpec("p", fn=process, loop=LoopSpec.range(0, 10, 2))
    graph = Graph.from_specs([spec])

    assert set(graph.specs.keys()) == {"p_0", "p_1", "p_2", "p_3", "p_4"}
    assert graph.specs["p_0"].args == (0,)
    assert graph.specs["p_4"].args == (8,)


def test_loop_range_two_args() -> None:
    """LoopSpec.range(start, stop) 两参数形式。"""

    def process(item: int) -> int:
        return item

    spec = TaskSpec("p", fn=process, loop=LoopSpec.range(5, 8))
    graph = Graph.from_specs([spec])

    assert set(graph.specs.keys()) == {"p_0", "p_1", "p_2"}
    assert graph.specs["p_0"].args == (5,)
    assert graph.specs["p_2"].args == (7,)


def test_loop_custom_key_fn() -> None:
    """key_fn 自定义实例名。"""

    def process(item: int) -> int:
        return item

    spec = TaskSpec(
        "process",
        fn=process,
        loop=LoopSpec.from_iterable([10, 20], key_fn=lambda _i, item: f"process_{item}"),
    )
    graph = Graph.from_specs([spec])

    assert set(graph.specs.keys()) == {"process_10", "process_20"}
    assert graph.specs["process_10"].args == (10,)


def test_loop_dependency_rewrite() -> None:
    """其他任务引用 loop 原名时，自动重写为依赖所有展开实例。"""

    def gen(item: int) -> int:
        return item

    def aggregate(gen: dict[str, int]) -> int:
        return sum(gen.values())

    gen_spec = TaskSpec("gen", fn=gen, loop=LoopSpec.range(3))
    agg_spec = TaskSpec("agg", fn=aggregate, depends_on=("gen",))
    graph = Graph.from_specs([gen_spec, agg_spec])

    assert set(graph.specs["agg"].depends_on) == {"gen_0", "gen_1", "gen_2"}


def test_loop_soft_dependency_rewrite() -> None:
    """软依赖引用 loop 原名时也重写。"""

    def gen(item: int) -> int:
        return item

    def consumer(gen: dict[str, int] | None = None) -> int:
        return sum(gen.values()) if gen else 0

    gen_spec = TaskSpec("gen", fn=gen, loop=LoopSpec.range(2))
    cons_spec = TaskSpec("cons", fn=consumer, soft_depends_on=("gen",))
    graph = Graph.from_specs([gen_spec, cons_spec])

    assert set(graph.specs["cons"].soft_depends_on) == {"gen_0", "gen_1"}


def test_loop_chained_dependencies() -> None:
    """两个 loop 任务链式依赖，下游 loop 引用上游 loop 原名。"""

    def stage1(item: int) -> int:
        return item + 1

    def stage2(stage1: int) -> int:
        return stage1 * 2

    s1 = TaskSpec("s1", fn=stage1, loop=LoopSpec.range(2))
    s2 = TaskSpec("s2", fn=stage2, loop=LoopSpec.range(2), depends_on=("s1",))
    graph = Graph.from_specs([s1, s2])

    # s2 的每个实例应依赖 s1 的所有实例（因为 s2 引用 "s1" 原名）
    assert set(graph.specs["s2_0"].depends_on) == {"s1_0", "s1_1"}
    assert set(graph.specs["s2_1"].depends_on) == {"s1_0", "s1_1"}


def test_loop_add_method() -> None:
    """Graph.add 也支持 loop 展开。"""

    def process(item: int) -> int:
        return item

    spec = TaskSpec("process", fn=process, loop=LoopSpec.range(2))
    graph = Graph().add(spec)

    assert set(graph.specs.keys()) == {"process_0", "process_1"}


def test_loop_execution() -> None:
    """loop 展开后能正常执行，结果注入下游。"""

    def gen(item: int) -> int:
        return item * 10

    def aggregate(gen_0: int, gen_1: int, gen_2: int) -> int:
        return gen_0 + gen_1 + gen_2

    gen_spec = TaskSpec("gen", fn=gen, loop=LoopSpec.range(3))
    agg_spec = TaskSpec("agg", fn=aggregate, depends_on=("gen",))
    graph = Graph.from_specs([gen_spec, agg_spec])

    report = px.run(graph, strategy="sequential")
    assert report.success
    assert report["agg"] == 0 + 10 + 20


def test_loop_with_task_decorator() -> None:
    """@task 装饰器支持 loop 参数。"""

    @task
    def base() -> int:
        return 0

    @task(loop=LoopSpec.range(2), depends_on=("base",))
    def worker(base: int, item: int) -> int:
        return base + item

    graph = Graph.from_specs([base, worker])
    assert set(graph.specs.keys()) == {"base", "worker_0", "worker_1"}
    assert set(graph.specs["worker_0"].depends_on) == {"base"}


def test_loop_groups_recorded() -> None:
    """_loop_groups 记录原名到展开名的映射。"""

    def process(item: int) -> int:
        return item

    spec = TaskSpec("p", fn=process, loop=LoopSpec.range(3))
    graph = Graph.from_specs([spec])

    assert graph._loop_groups == {"p": ["p_0", "p_1", "p_2"]}


def test_loop_empty_items_raises() -> None:
    """空 items 在 LoopSpec 构造时报错。"""
    import pytest

    with pytest.raises(ValueError, match="不能为空"):
        LoopSpec.from_iterable([])

    with pytest.raises(ValueError, match="不能为空"):
        LoopSpec.range(0)

    with pytest.raises(ValueError, match="不能为空"):
        LoopSpec.range(5, 5)  # start == stop, 空序列


def test_loop_preserves_other_fields() -> None:
    """loop 展开后保留原 spec 的其他字段（retry/tags 等）。"""

    def process(item: int) -> int:
        return item

    retry = px.RetryPolicy(max_attempts=3)
    spec = TaskSpec(
        "p",
        fn=process,
        loop=LoopSpec.range(2),
        retry=retry,
        tags=("batch",),
        priority=5,
    )
    graph = Graph.from_specs([spec])

    for name in ("p_0", "p_1"):
        s = graph.specs[name]
        assert s.retry == retry
        assert s.tags == ("batch",)
        assert s.priority == 5