stack/tests/aco/test_dag.py

"""Tests for aco.dag — DAG visualization from Pipeline introspection."""

from __future__ import annotations

import runpy
from unittest.mock import patch

import pytest

from aco.dag import (
    _EXTERNAL_COLOR,
    _SCHEMA_COLORS,
    Edge,
    Graph,
    Node,
    _color,
    _load_pipelines,
    build_graph,
    main,
    to_dot,
    to_html,
    to_mermaid,
)
from aco.express.base import Expr
from aco.pipe.base import Pipeline

# ── helpers ──────────────────────────────────────────────────────────────────


def _fn_a(core__encounter):
    """Compute A from encounter."""
    return core__encounter


def _fn_b(core__encounter, readmissions___int_encounter):
    """Compute B from encounter and int_encounter."""
    return core__encounter


def _fn_no_input():
    """A seed function with no inputs."""
    return None


def _make_expr(name, fn, description=""):
    return Expr(name=name, fn=fn, description=description)


def _make_pipeline(*exprs):
    return Pipeline(exprs=list(exprs))


# ── _color ───────────────────────────────────────────────────────────────────


class TestColor:
    def test_known_schema(self) -> None:
        assert _color("core") == _SCHEMA_COLORS["core"]
        assert _color("pharmacy") == _SCHEMA_COLORS["pharmacy"]

    def test_unknown_schema(self) -> None:
        assert _color("nonexistent_schema") == _EXTERNAL_COLOR


# ── data model ───────────────────────────────────────────────────────────────


class TestNode:
    def test_create(self) -> None:
        n = Node(name="core.encounter", schema="core", is_external=False)
        assert n.name == "core.encounter"
        assert n.schema == "core"
        assert n.is_external is False
        assert n.description == ""

    def test_create_with_description(self) -> None:
        n = Node(
            name="core.encounter",
            schema="core",
            is_external=False,
            description="An encounter",
        )
        assert n.description == "An encounter"


class TestEdge:
    def test_create(self) -> None:
        e = Edge(source="core.encounter", target="readmissions._int_encounter")
        assert e.source == "core.encounter"
        assert e.target == "readmissions._int_encounter"


class TestGraph:
    def test_empty(self) -> None:
        g = Graph()
        assert g.nodes == []
        assert g.edges == []

    def test_schemas(self) -> None:
        g = Graph(
            nodes=[
                Node("core.a", "core", False),
                Node("pharmacy.b", "pharmacy", False),
                Node("core.c", "core", True),
            ]
        )
        assert g.schemas() == {"core", "pharmacy"}

    def test_filter_none_returns_self(self) -> None:
        g = Graph(nodes=[Node("core.a", "core", False)])
        assert g.filter(None) is g

    def test_filter_with_schemas(self) -> None:
        n1 = Node("core.a", "core", False)
        n2 = Node("pharmacy.b", "pharmacy", False)
        n3 = Node("core.c", "core", False)
        g = Graph(
            nodes=[n1, n2, n3],
            edges=[
                Edge("core.a", "pharmacy.b"),
                Edge("core.a", "core.c"),
                Edge("pharmacy.b", "core.c"),
            ],
        )
        filtered = g.filter({"core"})
        assert len(filtered.nodes) == 2
        assert all(n.schema == "core" for n in filtered.nodes)
        # Only the edge between two core nodes survives
        assert len(filtered.edges) == 1
        assert filtered.edges[0].source == "core.a"
        assert filtered.edges[0].target == "core.c"


# ── build_graph ──────────────────────────────────────────────────────────────


class TestBuildGraph:
    def test_single_pipeline(self) -> None:
        expr_a = _make_expr("core.a", _fn_a, description="Expr A")
        p = _make_pipeline(expr_a)
        g = build_graph(p)
        names = {n.name for n in g.nodes}
        # core.a itself plus core.encounter as external input
        assert "core.a" in names
        assert "core.encounter" in names
        # The produced node is not external
        produced = [n for n in g.nodes if n.name == "core.a"][0]
        assert produced.is_external is False
        assert produced.description == "Expr A"
        # The input node is external
        ext = [n for n in g.nodes if n.name == "core.encounter"][0]
        assert ext.is_external is True

    def test_multiple_pipelines(self) -> None:
        expr_a = _make_expr("core.a", _fn_a)
        expr_b = _make_expr("readmissions.b", _fn_b)
        p1 = _make_pipeline(expr_a)
        p2 = _make_pipeline(expr_b)
        g = build_graph(p1, p2)
        names = {n.name for n in g.nodes}
        assert "core.a" in names
        assert "readmissions.b" in names

    def test_external_input_without_dot_uses_name_as_schema(self) -> None:
        # _fn_no_input has zero params, so nothing external.
        # We need a function with a param that has no "__" to become
        # an external input without "." in the name.
        def _fn_plain(seed):
            return seed

        expr = _make_expr("core.x", _fn_plain)
        p = _make_pipeline(expr)
        g = build_graph(p)
        ext = [n for n in g.nodes if n.name == "seed"][0]
        assert ext.is_external is True
        # When input has no dot, schema == the input name itself
        assert ext.schema == "seed"

    def test_edges(self) -> None:
        expr_a = _make_expr("core.a", _fn_a)
        p = _make_pipeline(expr_a)
        g = build_graph(p)
        assert any(
            e.source == "core.encounter" and e.target == "core.a" for e in g.edges
        )

    def test_produced_node_not_recreated_as_external(self) -> None:
        """If one expr produces a table another expr consumes, it's internal."""
        expr_enc = _make_expr("core.encounter", _fn_no_input)
        expr_a = _make_expr("core.a", _fn_a)
        p = _make_pipeline(expr_enc, expr_a)
        g = build_graph(p)
        enc_node = [n for n in g.nodes if n.name == "core.encounter"][0]
        assert enc_node.is_external is False


# ── to_dot ───────────────────────────────────────────────────────────────────


class TestToDot:
    @pytest.fixture()
    def graph(self) -> Graph:
        expr_a = _make_expr("core.a", _fn_a, description="Expr A")
        return build_graph(_make_pipeline(expr_a))

    def test_basic_structure(self, graph: Graph) -> None:
        dot = to_dot(graph)
        assert dot.startswith('digraph "ACO Pipeline DAG"')
        assert "rankdir=LR" in dot
        assert dot.rstrip().endswith("}")

    def test_custom_title(self, graph: Graph) -> None:
        dot = to_dot(graph, title="My Title")
        assert 'digraph "My Title"' in dot

    def test_subgraph_clusters(self, graph: Graph) -> None:
        dot = to_dot(graph)
        assert "subgraph cluster_core" in dot

    def test_external_node_dashed(self, graph: Graph) -> None:
        dot = to_dot(graph)
        # The external node core.encounter -> nid core__encounter
        assert 'style="dashed,filled"' in dot

    def test_internal_node_filled_rounded(self, graph: Graph) -> None:
        dot = to_dot(graph)
        assert 'style="filled,rounded"' in dot

    def test_edge_present(self, graph: Graph) -> None:
        dot = to_dot(graph)
        assert "core__encounter -> core__a;" in dot

    def test_node_without_dot_in_name(self) -> None:
        """Node whose name has no dot uses the full name as label."""

        def _fn_plain(seed):
            return seed

        expr = _make_expr("core.x", _fn_plain)
        g = build_graph(_make_pipeline(expr))
        dot = to_dot(g)
        # "seed" has no dot, so label is "seed" and nid is "seed"
        assert 'seed [label="seed"' in dot


# ── to_mermaid ───────────────────────────────────────────────────────────────


class TestToMermaid:
    @pytest.fixture()
    def graph(self) -> Graph:
        expr_a = _make_expr("core.a", _fn_a)
        return build_graph(_make_pipeline(expr_a))

    def test_header(self, graph: Graph) -> None:
        mmd = to_mermaid(graph)
        assert mmd.startswith("graph LR")

    def test_external_node_trapezoid(self, graph: Graph) -> None:
        mmd = to_mermaid(graph)
        # External nodes rendered as /name/
        assert "core_encounter[/encounter/]" in mmd

    def test_internal_node_rect(self, graph: Graph) -> None:
        mmd = to_mermaid(graph)
        assert "core_a[a]" in mmd

    def test_edge_arrow(self, graph: Graph) -> None:
        mmd = to_mermaid(graph)
        assert "core_encounter --> core_a" in mmd


# ── to_html ──────────────────────────────────────────────────────────────────


class TestToHtml:
    @pytest.fixture()
    def graph(self) -> Graph:
        expr_a = _make_expr("core.a", _fn_a, description="Expr A")
        return build_graph(_make_pipeline(expr_a))

    def test_doctype(self, graph: Graph) -> None:
        html = to_html(graph)
        assert "<!DOCTYPE html>" in html

    def test_default_title(self, graph: Graph) -> None:
        html = to_html(graph)
        assert "<title>ACO Pipeline DAG</title>" in html

    def test_custom_title(self, graph: Graph) -> None:
        html = to_html(graph, title="My <Title>")
        assert "My &lt;Title&gt;" in html

    def test_cytoscape_script(self, graph: Graph) -> None:
        html = to_html(graph)
        assert "cytoscape.min.js" in html
        assert "dagre.min.js" in html

    def test_elements_json_embedded(self, graph: Graph) -> None:
        html = to_html(graph)
        assert '"nodes"' in html
        assert '"edges"' in html

    def test_schema_colors_embedded(self, graph: Graph) -> None:
        html = to_html(graph)
        assert "schemaColors" in html
        assert _EXTERNAL_COLOR in html

    def test_external_node_in_elements(self, graph: Graph) -> None:
        html = to_html(graph)
        assert '"is_external": true' in html

    def test_group_nodes_present(self, graph: Graph) -> None:
        html = to_html(graph)
        assert '"is_group": true' in html
        assert "grp_core" in html

    def test_description_in_data(self, graph: Graph) -> None:
        html = to_html(graph)
        assert '"description": "Expr A"' in html


# ── _load_pipelines ─────────────────────────────────────────────────────────


class TestLoadPipelines:
    def test_specific_names(self) -> None:
        pipelines = _load_pipelines(["pharmacy"])
        assert len(pipelines) == 1
        assert isinstance(pipelines[0], Pipeline)
        assert len(pipelines[0].exprs) > 0

    def test_multiple_names(self) -> None:
        pipelines = _load_pipelines(["pharmacy", "readmissions"])
        assert len(pipelines) == 2

    def test_none_loads_all(self) -> None:
        pipelines = _load_pipelines(None)
        assert len(pipelines) == 13  # _ALL_PIPE_MODULES has 13 entries


# ── main CLI ─────────────────────────────────────────────────────────────────


class TestMain:
    """Test CLI entry point by mocking sys.argv and capturing output."""

    def _make_mock_pipeline(self):
        """Return a Pipeline with a simple expr for fast CLI tests."""
        expr = _make_expr("core.a", _fn_a, description="test")
        return _make_pipeline(expr)

    def _patch_load(self):
        """Patch _load_pipelines to return a tiny pipeline."""
        p = self._make_mock_pipeline()
        return patch("aco.dag._load_pipelines", return_value=[p])

    def test_summary_mode(self, capsys) -> None:
        with self._patch_load(), patch("sys.argv", ["dag"]):
            main()
        out = capsys.readouterr().out
        assert "Nodes:" in out
        assert "Edges:" in out
        assert "Schemas:" in out

    def test_html_to_file(self, tmp_path, capsys) -> None:
        outfile = str(tmp_path / "dag.html")
        with self._patch_load(), patch("sys.argv", ["dag", "--html", outfile]):
            main()
        out = capsys.readouterr().out
        assert f"Wrote {outfile}" in out
        content = (tmp_path / "dag.html").read_text()
        assert "<!DOCTYPE html>" in content

    def test_html_to_stdout(self, capsys) -> None:
        with self._patch_load(), patch("sys.argv", ["dag", "--html", "-"]):
            main()
        out = capsys.readouterr().out
        assert "<!DOCTYPE html>" in out

    def test_html_with_pipeline_flag(self, capsys) -> None:
        with (
            self._patch_load(),
            patch("sys.argv", ["dag", "-p", "pharmacy", "--html", "-"]),
        ):
            main()
        out = capsys.readouterr().out
        assert "pharmacy DAG" in out

    def test_dot_to_file(self, tmp_path, capsys) -> None:
        outfile = str(tmp_path / "dag.dot")
        with self._patch_load(), patch("sys.argv", ["dag", "--dot", outfile]):
            main()
        out = capsys.readouterr().out
        assert f"Wrote {outfile}" in out
        content = (tmp_path / "dag.dot").read_text()
        assert "digraph" in content

    def test_dot_to_stdout(self, capsys) -> None:
        with self._patch_load(), patch("sys.argv", ["dag", "--dot", "-"]):
            main()
        out = capsys.readouterr().out
        assert "digraph" in out

    def test_mermaid_mode(self, capsys) -> None:
        with self._patch_load(), patch("sys.argv", ["dag", "--mermaid"]):
            main()
        out = capsys.readouterr().out
        assert "graph LR" in out

    def test_module_main_guard(self, capsys) -> None:
        """Cover the ``if __name__ == '__main__'`` block via runpy."""
        with self._patch_load(), patch("sys.argv", ["dag"]):
            runpy.run_module("aco.dag", run_name="__main__")