hokusai/tests/test_graph.py

"""Integration tests for hokusai.graph."""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path

import pytest

from hokusai.config import ProjectConfig
from hokusai.graph import build_graph, get_build_order, get_subgraph_for_target

WriteConfig = Callable[[dict[str, object]], ProjectConfig]


class TestBuildGraph:
    """Test dependency graph construction from real configs."""

    def test_single_target_no_deps(
        self, project_dir: Path, simple_text_config: ProjectConfig
    ) -> None:
        graph = build_graph(simple_text_config, project_dir)
        assert "output.txt" in graph.nodes
        assert graph.number_of_edges() == 0

    def test_chain_dependency(
        self, project_dir: Path, multi_target_config: ProjectConfig
    ) -> None:
        graph = build_graph(multi_target_config, project_dir)

        assert graph.has_edge("input.txt", "summary.md")
        assert graph.has_edge("summary.md", "final.txt")
        assert not graph.has_edge("input.txt", "final.txt")

    def test_external_file_as_node(
        self, project_dir: Path, multi_target_config: ProjectConfig
    ) -> None:
        graph = build_graph(multi_target_config, project_dir)
        assert "input.txt" in graph.nodes

    def test_missing_dependency_raises(
        self, project_dir: Path, write_config: WriteConfig
    ) -> None:
        config = write_config(
            {"targets": {"out.txt": {"prompt": "x", "inputs": ["nonexistent.txt"]}}}
        )
        with pytest.raises(
            ValueError, match="neither a defined target nor an existing file"
        ):
            _ = build_graph(config, project_dir)

    def test_cycle_raises(self, project_dir: Path, write_config: WriteConfig) -> None:
        config = write_config(
            {
                "targets": {
                    "a.txt": {"prompt": "x", "inputs": ["b.txt"]},
                    "b.txt": {"prompt": "x", "inputs": ["a.txt"]},
                }
            }
        )
        with pytest.raises(ValueError, match="cycle"):
            _ = build_graph(config, project_dir)

    def test_reference_image_creates_edge(
        self, project_dir: Path, write_config: WriteConfig
    ) -> None:
        _ = (project_dir / "ref.png").write_bytes(b"\x89PNG")
        config = write_config(
            {"targets": {"out.png": {"prompt": "x", "reference_images": ["ref.png"]}}}
        )
        graph = build_graph(config, project_dir)
        assert graph.has_edge("ref.png", "out.png")

    def test_control_images_create_edges(
        self, project_dir: Path, write_config: WriteConfig
    ) -> None:
        _ = (project_dir / "ctrl1.png").write_bytes(b"\x89PNG")
        _ = (project_dir / "ctrl2.png").write_bytes(b"\x89PNG")
        config = write_config(
            {
                "targets": {
                    "out.png": {
                        "prompt": "x",
                        "control_images": ["ctrl1.png", "ctrl2.png"],
                    }
                }
            }
        )
        graph = build_graph(config, project_dir)
        assert graph.has_edge("ctrl1.png", "out.png")
        assert graph.has_edge("ctrl2.png", "out.png")

    def test_target_depending_on_another_target(
        self, project_dir: Path, write_config: WriteConfig
    ) -> None:
        config = write_config(
            {
                "targets": {
                    "base.txt": {"prompt": "base"},
                    "derived.txt": {"prompt": "derive", "inputs": ["base.txt"]},
                }
            }
        )
        graph = build_graph(config, project_dir)
        assert graph.has_edge("base.txt", "derived.txt")


class TestGetBuildOrder:
    """Test topological generation ordering."""

    def test_independent_targets_same_generation(
        self, project_dir: Path, write_config: WriteConfig
    ) -> None:
        config = write_config(
            {
                "targets": {
                    "a.txt": {"prompt": "x"},
                    "b.txt": {"prompt": "y"},
                    "c.png": {"prompt": "z"},
                }
            }
        )
        graph = build_graph(config, project_dir)
        order = get_build_order(graph)

        assert len(order) == 1
        assert set(order[0]) == {"a.txt", "b.txt", "c.png"}

    def test_chain_produces_sequential_generations(
        self, project_dir: Path, multi_target_config: ProjectConfig
    ) -> None:
        graph = build_graph(multi_target_config, project_dir)
        order = get_build_order(graph)

        # Flatten to find relative positions
        flat = [name for gen in order for name in gen]
        assert flat.index("input.txt") < flat.index("summary.md")
        assert flat.index("summary.md") < flat.index("final.txt")

    def test_diamond_dependency(
        self, project_dir: Path, write_config: WriteConfig
    ) -> None:
        _ = (project_dir / "root.txt").write_text("root")
        config = write_config(
            {
                "targets": {
                    "left.txt": {"prompt": "x", "inputs": ["root.txt"]},
                    "right.txt": {"prompt": "y", "inputs": ["root.txt"]},
                    "merge.txt": {
                        "prompt": "z",
                        "inputs": ["left.txt", "right.txt"],
                    },
                }
            }
        )
        graph = build_graph(config, project_dir)
        order = get_build_order(graph)
        flat = [name for gen in order for name in gen]

        assert flat.index("root.txt") < flat.index("left.txt")
        assert flat.index("root.txt") < flat.index("right.txt")
        assert flat.index("left.txt") < flat.index("merge.txt")
        assert flat.index("right.txt") < flat.index("merge.txt")


class TestGetSubgraphForTarget:
    """Test selective subgraph extraction."""

    def test_subgraph_includes_transitive_deps(
        self, project_dir: Path, multi_target_config: ProjectConfig
    ) -> None:
        graph = build_graph(multi_target_config, project_dir)
        sub = get_subgraph_for_target(graph, "final.txt")

        assert "final.txt" in sub.nodes
        assert "summary.md" in sub.nodes
        assert "input.txt" in sub.nodes
        # hero.png is independent and should NOT be included
        assert "hero.png" not in sub.nodes

    def test_subgraph_leaf_target(
        self, project_dir: Path, multi_target_config: ProjectConfig
    ) -> None:
        graph = build_graph(multi_target_config, project_dir)
        sub = get_subgraph_for_target(graph, "hero.png")

        assert set(sub.nodes) == {"hero.png"}

    def test_subgraph_preserves_edges(
        self, project_dir: Path, multi_target_config: ProjectConfig
    ) -> None:
        graph = build_graph(multi_target_config, project_dir)
        sub = get_subgraph_for_target(graph, "final.txt")

        assert sub.has_edge("input.txt", "summary.md")
        assert sub.has_edge("summary.md", "final.txt")