pyflowx/tests/test_group.py

"""任务分组（Graph.group）测试。

覆盖：
* 基本分组创建与依赖重写（硬依赖 / 软依赖）。
* 组与 loop 组合：组内成员可包含 loop 展开后的实例名。
* 校验：空组名、与任务名冲突、组名重复、组内任务未注册、空组列表。
* 链式调用：group() 返回 self。
* 执行：分组后图能正常跑通。
"""

from __future__ import annotations

import pytest

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


def test_group_basic_creation() -> None:
    """group() 记录到 _groups 字典并返回 self。"""
    graph = Graph.from_specs(
        [
            TaskSpec("a", fn=lambda: 1),
            TaskSpec("b", fn=lambda: 2),
        ]
    )
    result = graph.group("phase1", ["a", "b"])
    assert result is graph
    assert graph._groups["phase1"] == ("a", "b")


def test_group_dependency_rewrite_hard() -> None:
    """引用组名的硬依赖被重写为组内全部任务。"""
    graph = Graph.from_specs(
        [
            TaskSpec("a", fn=lambda: 1),
            TaskSpec("b", fn=lambda: 2),
        ]
    )
    graph.group("phase1", ["a", "b"])
    graph.add(TaskSpec("c", fn=lambda: 3, depends_on=("phase1",)))
    assert graph.specs["c"].depends_on == ("a", "b")


def test_group_dependency_rewrite_soft() -> None:
    """引用组名的软依赖被重写为组内全部任务。"""
    graph = Graph.from_specs(
        [
            TaskSpec("a", fn=lambda: 1),
            TaskSpec("b", fn=lambda: 2),
        ]
    )
    graph.group("phase1", ["a", "b"])
    graph.add(TaskSpec("c", fn=lambda: 3, soft_depends_on=("phase1",)))
    assert graph.specs["c"].soft_depends_on == ("a", "b")


def test_group_before_adding_depender() -> None:
    """先声明 group 再注册引用方：from_specs 后 add 引用组名。"""
    graph = Graph.from_specs(
        [
            TaskSpec("extract", fn=lambda: [1, 2]),
            TaskSpec("validate", fn=lambda: True),
        ]
    )
    graph.group("prep", ["extract", "validate"])
    graph.add(TaskSpec("process", fn=lambda: None, depends_on=("prep",)))
    assert set(graph.specs["process"].depends_on) == {"extract", "validate"}


def test_group_chained_with_loop() -> None:
    """组内成员可包含 loop 展开后的实例名。"""
    graph = Graph.from_specs(
        [
            TaskSpec("proc", fn=lambda x: x * 2, loop=LoopSpec.range(3)),
        ]
    )
    # proc 展开为 proc_0, proc_1, proc_2
    assert {"proc_0", "proc_1", "proc_2"}.issubset(graph.specs.keys())
    graph.group("proc_group", ["proc_0", "proc_1", "proc_2"])
    graph.add(TaskSpec("post", fn=lambda: None, depends_on=("proc_group",)))
    assert set(graph.specs["post"].depends_on) == {"proc_0", "proc_1", "proc_2"}


def test_group_execution() -> None:
    """分组后的图能正常执行。"""
    calls: list[str] = []

    def make(name: str):
        def _f() -> None:
            calls.append(name)

        return _f

    graph = Graph.from_specs(
        [
            TaskSpec("a", fn=make("a")),
            TaskSpec("b", fn=make("b")),
            TaskSpec("c", fn=make("c")),
        ]
    )
    graph.group("prep", ["a", "b"])
    graph.add(TaskSpec("done", fn=make("done"), depends_on=("prep", "c")))
    report = px.run(graph, strategy="sequential")
    assert report.success
    assert set(calls) == {"a", "b", "c", "done"}
    # done 在 a/b/c 之后
    assert calls.index("done") > calls.index("a")
    assert calls.index("done") > calls.index("b")
    assert calls.index("done") > calls.index("c")


def test_group_empty_name_raises() -> None:
    """空组名抛 ValueError。"""
    graph = Graph.from_specs([TaskSpec("a", fn=lambda: 1)])
    with pytest.raises(ValueError, match="非空字符串"):
        graph.group("", ["a"])


def test_group_name_conflicts_with_task() -> None:
    """组名与已注册任务名冲突时抛 ValueError。"""
    graph = Graph.from_specs([TaskSpec("a", fn=lambda: 1)])
    with pytest.raises(ValueError, match="与已注册任务名冲突"):
        graph.group("a", [])


def test_group_duplicate_name_raises() -> None:
    """重复声明同名组抛 ValueError。"""
    graph = Graph.from_specs(
        [
            TaskSpec("a", fn=lambda: 1),
            TaskSpec("b", fn=lambda: 2),
        ]
    )
    graph.group("g1", ["a"])
    with pytest.raises(ValueError, match="已存在"):
        graph.group("g1", ["b"])


def test_group_unregistered_task_raises() -> None:
    """组内引用未注册的任务抛 ValueError。"""
    graph = Graph.from_specs([TaskSpec("a", fn=lambda: 1)])
    with pytest.raises(ValueError, match="未注册的任务"):
        graph.group("g1", ["a", "nonexistent"])


def test_group_empty_tasks_raises() -> None:
    """空组列表抛 ValueError。"""
    graph = Graph.from_specs([TaskSpec("a", fn=lambda: 1)])
    with pytest.raises(ValueError, match="不能为空"):
        graph.group("g1", [])


def test_group_chain_returns_self() -> None:
    """group() 返回 self 支持链式调用。"""
    graph = Graph.from_specs(
        [
            TaskSpec("a", fn=lambda: 1),
            TaskSpec("b", fn=lambda: 2),
            TaskSpec("c", fn=lambda: 3),
        ]
    )
    result = graph.group("g1", ["a"]).group("g2", ["b", "c"])
    assert result is graph
    assert graph._groups["g1"] == ("a",)
    assert graph._groups["g2"] == ("b", "c")


def test_group_mixed_hard_soft() -> None:
    """同一任务的硬依赖和软依赖可同时引用不同组。"""
    graph = Graph.from_specs(
        [
            TaskSpec("a", fn=lambda: 1),
            TaskSpec("b", fn=lambda: 2),
            TaskSpec("c", fn=lambda: 3),
            TaskSpec("d", fn=lambda: 4),
        ]
    )
    graph.group("g1", ["a", "b"])
    graph.group("g2", ["c", "d"])
    graph.add(
        TaskSpec(
            "e",
            fn=lambda: 5,
            depends_on=("g1",),
            soft_depends_on=("g2",),
        )
    )
    assert set(graph.specs["e"].depends_on) == {"a", "b"}
    assert set(graph.specs["e"].soft_depends_on) == {"c", "d"}