from __future__ import annotations import unittest from sqlalchemy import create_engine, func, select from sqlalchemy.orm import Session, sessionmaker from govoplan_core.auth import ApiPrincipal from govoplan_core.core.access import PrincipalRef from govoplan_core.core.dataflows import ( DataflowPublicationTarget, DataflowRunRequest, ) from govoplan_core.core.datasources import ( CAPABILITY_DATASOURCE_PUBLICATION, DatasourceDescriptor, DatasourceMaterialization, DatasourcePublicationResult, ) from govoplan_core.db.base import Base from govoplan_dataflow.backend.db.models import ( DataflowPipeline, DataflowPipelineRevision, DataflowRun, ) from govoplan_dataflow.backend.schemas import ( GraphEdge, GraphNode, GraphPosition, PipelineCreateRequest, PipelineGraph, PipelinePreviewRequest, PipelineUpdateRequest, ) from govoplan_dataflow.backend.service import ( DataflowConflictError, DataflowNotFoundError, create_pipeline, delete_pipeline, get_pipeline, list_pipelines, preview_pipeline, start_pipeline_run, update_pipeline, ) def sample_graph(*, minimum: int = 10) -> PipelineGraph: return PipelineGraph( nodes=[ GraphNode( id="source", type="source.inline", label="Monthly input", position=GraphPosition(x=40, y=160), config={ "source_name": "monthly_files", "rows": [ {"id": 1, "amount": 5}, {"id": 2, "amount": 15}, {"id": 3, "amount": 25}, ], }, ), GraphNode( id="filter", type="filter", label="Minimum amount", position=GraphPosition(x=260, y=160), config={"column": "amount", "operator": "gte", "value": minimum}, ), GraphNode( id="output", type="output", label="Output", position=GraphPosition(x=480, y=160), config={}, ), ], edges=[ GraphEdge(id="source-filter", source="source", target="filter"), GraphEdge(id="filter-output", source="filter", target="output"), ], ) def principal(tenant_id: str = "tenant-1") -> ApiPrincipal: return ApiPrincipal( principal=PrincipalRef( account_id="account-1", membership_id="membership-1", tenant_id=tenant_id, scopes=frozenset(), ), account=object(), user=object(), ) class FakePublicationProvider: def __init__(self) -> None: self.requests = [] def publish_rows(self, _session, _principal, *, request): self.requests.append(request) descriptor = DatasourceDescriptor( ref=request.target_datasource_ref or "datasource:output-1", source_name=request.source_name or "existing_output", name=request.name or "Existing output", kind="custom", mode="static", shape="tabular", fingerprint="published-fingerprint", ) return DatasourcePublicationResult( ref="publication:publication-1", status="published", datasource=descriptor, materialization=DatasourceMaterialization( ref="materialization:materialization-1", datasource_ref=descriptor.ref, revision=1, state="published", fingerprint=descriptor.fingerprint, ), ) class FakeRegistry: def __init__(self, publication_provider: FakePublicationProvider) -> None: self.publication_provider = publication_provider def has_capability(self, name: str) -> bool: return name == CAPABILITY_DATASOURCE_PUBLICATION def capability(self, name: str): if not self.has_capability(name): raise KeyError(name) return self.publication_provider class DataflowServiceTests(unittest.TestCase): def setUp(self) -> None: self.engine = create_engine("sqlite:///:memory:") Base.metadata.create_all( self.engine, tables=[ DataflowPipeline.__table__, DataflowPipelineRevision.__table__, DataflowRun.__table__, ], ) self.Session = sessionmaker(bind=self.engine) self.session: Session = self.Session() def tearDown(self) -> None: self.session.close() Base.metadata.drop_all( self.engine, tables=[ DataflowRun.__table__, DataflowPipelineRevision.__table__, DataflowPipeline.__table__, ], ) self.engine.dispose() def _create(self, *, tenant_id: str = "tenant-1") -> DataflowPipeline: pipeline = create_pipeline( self.session, tenant_id=tenant_id, actor_id="user-1", payload=PipelineCreateRequest( name="Monthly comparison", description="First governed pipeline", status="active", graph=sample_graph(), editor_mode="graph", ), ) self.session.commit() return pipeline def test_create_and_update_make_immutable_revisions(self) -> None: pipeline = self._create() updated = update_pipeline( self.session, tenant_id="tenant-1", pipeline_id=pipeline.id, actor_id="user-2", payload=PipelineUpdateRequest( name="Monthly comparison", description="Changed threshold", graph=sample_graph(minimum=20), editor_mode="graph", expected_revision=1, ), ) self.session.commit() revisions = list( self.session.scalars( select(DataflowPipelineRevision) .where(DataflowPipelineRevision.pipeline_id == pipeline.id) .order_by(DataflowPipelineRevision.revision) ) ) self.assertEqual(2, updated.current_revision) self.assertEqual([1, 2], [item.revision for item in revisions]) self.assertNotEqual(revisions[0].content_hash, revisions[1].content_hash) self.assertEqual(10, revisions[0].graph["nodes"][1]["config"]["value"]) self.assertEqual(20, revisions[1].graph["nodes"][1]["config"]["value"]) def test_stale_revision_is_rejected(self) -> None: pipeline = self._create() with self.assertRaises(DataflowConflictError): update_pipeline( self.session, tenant_id="tenant-1", pipeline_id=pipeline.id, actor_id="user-2", payload=PipelineUpdateRequest( name="Stale edit", graph=sample_graph(minimum=30), editor_mode="graph", expected_revision=2, ), ) def test_pipeline_lookup_is_tenant_isolated(self) -> None: pipeline = self._create() with self.assertRaises(DataflowNotFoundError): get_pipeline( self.session, tenant_id="tenant-2", pipeline_id=pipeline.id, ) def test_list_and_soft_delete_are_tenant_isolated(self) -> None: first = self._create() second = self._create(tenant_id="tenant-2") self.assertEqual([first.id], [item.id for item in list_pipelines(self.session, tenant_id="tenant-1")]) self.assertEqual([second.id], [item.id for item in list_pipelines(self.session, tenant_id="tenant-2")]) delete_pipeline( self.session, tenant_id="tenant-1", pipeline_id=first.id, actor_id="user-2", ) self.session.commit() self.assertEqual([], list_pipelines(self.session, tenant_id="tenant-1")) with self.assertRaises(DataflowNotFoundError): get_pipeline( self.session, tenant_id="tenant-1", pipeline_id=first.id, ) self.assertEqual( 1, self.session.scalar( select(func.count()) .select_from(DataflowPipelineRevision) .where(DataflowPipelineRevision.pipeline_id == first.id) ), ) def test_saved_preview_records_lineage_but_not_result_rows(self) -> None: pipeline = self._create() response = preview_pipeline( self.session, tenant_id="tenant-1", actor_id="user-1", payload=PipelinePreviewRequest( pipeline_id=pipeline.id, preview_node_id="source", row_limit=1, ), principal=principal(), ) self.session.commit() run = self.session.scalar(select(DataflowRun).where(DataflowRun.id == response.run_id)) self.assertEqual("succeeded", response.status) self.assertEqual([{"id": 2, "amount": 15}], response.rows) self.assertEqual(2, response.total_rows) self.assertTrue(response.truncated) self.assertIsNotNone(response.node_preview) self.assertEqual("source", response.node_preview.node_id) self.assertEqual([{"id": 1, "amount": 5}], response.node_preview.rows) self.assertEqual(3, response.node_preview.total_rows) self.assertTrue(response.node_preview.truncated) self.assertEqual(2, run.output_row_count) self.assertEqual(3, run.input_row_count) self.assertEqual(1, len(run.source_fingerprints)) self.assertEqual(3, response.input_row_count) self.assertEqual(run.source_fingerprints, response.source_fingerprints) self.assertFalse(hasattr(run, "result_rows")) self.assertEqual( 1, self.session.scalar(select(func.count()).select_from(DataflowRun)), ) def test_failed_preview_keeps_upstream_lineage_and_failed_node_diagnostic(self) -> None: graph = sample_graph() graph.nodes[0].config["rows"] = [{"id": 1, "amount": "not-a-number"}] pipeline = create_pipeline( self.session, tenant_id="tenant-1", actor_id="user-1", payload=PipelineCreateRequest( name="Invalid runtime value", graph=graph, editor_mode="graph", ), ) self.session.commit() response = preview_pipeline( self.session, tenant_id="tenant-1", actor_id="user-1", payload=PipelinePreviewRequest(pipeline_id=pipeline.id), principal=principal(), ) self.session.commit() run = self.session.scalar(select(DataflowRun).where(DataflowRun.id == response.run_id)) self.assertEqual("failed", response.status) self.assertEqual( [("source", "succeeded"), ("filter", "failed")], [(item.node_id, item.status) for item in response.node_diagnostics], ) self.assertEqual(1, run.input_row_count) self.assertEqual(1, len(run.source_fingerprints)) self.assertEqual(1, response.input_row_count) self.assertEqual(run.source_fingerprints, response.source_fingerprints) self.assertEqual("preview.execution", run.diagnostics[-1]["code"]) def test_pinned_run_publishes_once_and_replays_idempotently(self) -> None: pipeline = self._create() publication_provider = FakePublicationProvider() request = DataflowRunRequest( pipeline_ref=f"pipeline:{pipeline.id}", revision=1, idempotency_key="monthly-2026-07", publication=DataflowPublicationTarget( name="Monthly result", source_name="monthly_result", freeze=True, frozen_label="July 2026", ), ) first, first_replayed = start_pipeline_run( self.session, tenant_id="tenant-1", actor_id="user-1", principal=principal(), registry=FakeRegistry(publication_provider), request=request, ) replay, replayed = start_pipeline_run( self.session, tenant_id="tenant-1", actor_id="user-1", principal=principal(), registry=FakeRegistry(publication_provider), request=request, ) self.session.commit() self.assertFalse(first_replayed) self.assertTrue(replayed) self.assertEqual(first.id, replay.id) self.assertEqual("succeeded", first.status) self.assertEqual("datasource:output-1", first.output_datasource_ref) self.assertEqual( "materialization:materialization-1", first.output_materialization_ref, ) self.assertEqual( [{"id": 2, "amount": 15}, {"id": 3, "amount": 25}], list(publication_provider.requests[0].rows), ) self.assertEqual(1, len(publication_provider.requests)) def test_run_idempotency_key_rejects_changed_parameters(self) -> None: pipeline = self._create() publication_provider = FakePublicationProvider() registry = FakeRegistry(publication_provider) start_pipeline_run( self.session, tenant_id="tenant-1", actor_id="user-1", principal=principal(), registry=registry, request=DataflowRunRequest( pipeline_ref=f"pipeline:{pipeline.id}", revision=1, idempotency_key="stable-key", row_limit=100, ), ) with self.assertRaises(DataflowConflictError): start_pipeline_run( self.session, tenant_id="tenant-1", actor_id="user-1", principal=principal(), registry=registry, request=DataflowRunRequest( pipeline_ref=f"pipeline:{pipeline.id}", revision=1, idempotency_key="stable-key", row_limit=10, ), ) def test_publication_without_datasources_finishes_as_failed_run(self) -> None: pipeline = self._create() run, _ = start_pipeline_run( self.session, tenant_id="tenant-1", actor_id="user-1", principal=principal(), registry=None, request=DataflowRunRequest( pipeline_ref=f"pipeline:{pipeline.id}", revision=1, idempotency_key="missing-publisher", publication=DataflowPublicationTarget( name="Output", source_name="output", ), ), ) self.assertEqual("failed", run.status) self.assertIn("Datasources publication capability", run.error) if __name__ == "__main__": unittest.main()