govoplan-dataflow/tests/test_service.py

from __future__ import annotations

import unittest

from sqlalchemy import create_engine, func, select
from sqlalchemy.orm import Session, sessionmaker

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,
    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"),
        ],
    )


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",
                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, row_limit=1),
        )
        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.assertEqual(2, run.output_row_count)
        self.assertEqual(3, run.input_row_count)
        self.assertEqual(1, len(run.source_fingerprints))
        self.assertFalse(hasattr(run, "result_rows"))
        self.assertEqual(
            1,
            self.session.scalar(select(func.count()).select_from(DataflowRun)),
        )


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