govoplan-dataflow/tests/test_golden_flows.py

from __future__ import annotations

import json
import hashlib
import unittest
from pathlib import Path

from govoplan_dataflow.backend.executor import (
    PipelineExecutionError,
    ResolvedSource,
    execute_preview,
)
from govoplan_dataflow.backend.expressions import ExpressionError, parse_expression
from govoplan_dataflow.backend.graph import definition_hash, validate_graph
from govoplan_dataflow.backend.schemas import PipelineGraph
from govoplan_dataflow.backend.subflows import substitute_parameters


FIXTURES = Path(__file__).parents[1] / "fixtures" / "golden"


class GoldenFlowTests(unittest.TestCase):
    def test_every_golden_flow_matches_its_expected_output(self) -> None:
        directories = sorted(
            path for path in FIXTURES.iterdir() if path.is_dir()
        )
        self.assertGreaterEqual(len(directories), 2)
        for directory in directories:
            with self.subTest(flow=directory.name):
                manifest = _json(directory / "manifest.json")
                raw_graph_payload = _json(directory / str(manifest["graph"]))
                graph_payload = json.loads(json.dumps(raw_graph_payload))
                for item in graph_payload["nodes"]:
                    fixture = item.get("config", {}).get("fixture")
                    if fixture and item["type"] == "source.inline":
                        item["config"]["rows"] = _json(
                            directory / "inputs" / fixture
                        )
                graph_payload = substitute_parameters(
                    graph_payload,
                    dict(manifest.get("parameters") or {}),
                )
                graph = PipelineGraph.model_validate(graph_payload)
                diagnostics = validate_graph(graph)
                self.assertEqual(
                    [],
                    [
                        item.model_dump()
                        for item in diagnostics
                        if item.severity == "error"
                    ],
                )

                resolver = _source_resolver(directory, manifest)
                result = execute_preview(
                    graph,
                    row_limit=500,
                    source_resolver=resolver,
                )

                self.assertEqual(
                    _json(directory / str(manifest["expected_output"])),
                    result.rows,
                )
                if "expected_definition_hash" in manifest:
                    self.assertEqual(
                        manifest["expected_definition_hash"],
                        definition_hash(PipelineGraph.model_validate(raw_graph_payload)),
                    )
                if "expected_execution_graph_hash" in manifest:
                    self.assertEqual(
                        manifest["expected_execution_graph_hash"],
                        definition_hash(graph),
                    )
                if "expected_output_hash" in manifest:
                    self.assertEqual(
                        manifest["expected_output_hash"],
                        _rows_hash(result.rows),
                    )
                if "expected_source_fingerprints" in manifest:
                    self.assertEqual(
                        manifest["expected_source_fingerprints"],
                        result.source_fingerprints,
                    )
                full_result = execute_preview(
                    graph,
                    row_limit=5_000,
                    source_resolver=resolver,
                )
                self.assertEqual(result.rows, full_result.rows)
                self.assertFalse(full_result.truncated)

    def test_adrema_fixture_rejects_untrusted_sources_and_unsafe_expression(self) -> None:
        directory = FIXTURES / "adrema-audience-selection"
        manifest = _json(directory / "manifest.json")
        graph = PipelineGraph.model_validate(
            substitute_parameters(
                _json(directory / "graph.json"),
                dict(manifest["parameters"]),
            )
        )
        drifted_manifest = {
            **manifest,
            "sources": {
                **manifest["sources"],
                "snapshot:directory-2026-01": {
                    **manifest["sources"]["snapshot:directory-2026-01"],
                    "fingerprint": "sha256:changed-source",
                },
            },
        }

        with self.assertRaisesRegex(PipelineExecutionError, "fingerprint"):
            execute_preview(
                graph,
                row_limit=500,
                source_resolver=_source_resolver(directory, drifted_manifest),
            )
        cross_tenant_manifest = {
            **manifest,
            "sources": {
                **manifest["sources"],
                "snapshot:directory-2026-01": {
                    **manifest["sources"]["snapshot:directory-2026-01"],
                    "tenant_id": "tenant-other",
                },
            },
        }
        with self.assertRaisesRegex(PipelineExecutionError, "tenant"):
            execute_preview(
                graph,
                row_limit=500,
                source_resolver=_source_resolver(directory, cross_tenant_manifest),
            )

        unknown_source_graph = graph.model_copy(deep=True)
        unknown_source_graph.nodes[0].config["source_ref"] = "connector:untrusted"
        with self.assertRaisesRegex(PipelineExecutionError, "pinned provider set"):
            execute_preview(
                unknown_source_graph,
                row_limit=500,
                source_resolver=_source_resolver(directory, manifest),
            )
        with self.assertRaises(ExpressionError):
            parse_expression("read_csv('/tmp/private.csv')")


def _json(path: Path):
    return json.loads(path.read_text(encoding="utf-8"))


def _source_resolver(directory: Path, manifest: dict):
    sources = dict(manifest.get("sources") or {})
    tenant_id = str(manifest.get("tenant_id") or "")

    def resolve(node, limit: int) -> ResolvedSource:
        source_ref = str(node.config.get("source_ref") or "")
        if source_ref not in sources:
            raise ValueError(f"Fixture source {source_ref!r} is outside the pinned provider set.")
        source = sources[source_ref]
        if str(source.get("tenant_id") or "") != tenant_id:
            raise ValueError(
                f"Fixture source {source_ref!r} belongs to a different tenant."
            )
        fingerprint = str(source["fingerprint"])
        expected = str(node.config.get("expected_fingerprint") or "")
        if expected and expected != fingerprint:
            raise ValueError(
                f"Fixture source fingerprint changed for {source_ref!r}."
            )
        rows = _json(directory / str(source["fixture"]))
        bounded = rows[:limit]
        return ResolvedSource(
            rows=tuple(dict(item) for item in bounded),
            source_ref=source_ref,
            provider=str(source["provider"]),
            fingerprint=fingerprint,
            total_rows=len(rows),
            truncated=len(bounded) < len(rows),
        )

    return resolve


def _rows_hash(rows: object) -> str:
    return hashlib.sha256(
        json.dumps(
            rows,
            sort_keys=True,
            separators=(",", ":"),
            ensure_ascii=True,
        ).encode("utf-8")
    ).hexdigest()


if __name__ == "__main__":
    unittest.main()