Add governed Dataflow audience outputs
This commit is contained in:
+147
-5
@@ -1,12 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import hashlib
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from govoplan_dataflow.backend.executor import execute_preview
|
||||
from govoplan_dataflow.backend.graph import validate_graph
|
||||
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"
|
||||
@@ -21,13 +28,18 @@ class GoldenFlowTests(unittest.TestCase):
|
||||
for directory in directories:
|
||||
with self.subTest(flow=directory.name):
|
||||
manifest = _json(directory / "manifest.json")
|
||||
graph_payload = _json(directory / str(manifest["graph"]))
|
||||
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:
|
||||
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(
|
||||
@@ -39,17 +51,147 @@ class GoldenFlowTests(unittest.TestCase):
|
||||
],
|
||||
)
|
||||
|
||||
result = execute_preview(graph, row_limit=500)
|
||||
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()
|
||||
|
||||
Reference in New Issue
Block a user