govoplan-dataflow/tests/test_run_worker.py

from __future__ import annotations

from datetime import timedelta
import unittest

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

from govoplan_core.auth import ApiPrincipal
from govoplan_core.core.access import (
    CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER,
    PrincipalRef,
)
from govoplan_core.core.automation import AutomationPrincipalResolution
from govoplan_core.core.dataflows import DataflowRunRequest
from govoplan_core.db.base import Base, utcnow
from govoplan_dataflow.backend.db.models import (
    DataflowPipeline,
    DataflowPipelineDeployment,
    DataflowPipelineRevision,
    DataflowRun,
)
from govoplan_dataflow.backend.run_worker import (
    dispatch_pending_runs,
    purge_expired_runs,
)
from govoplan_dataflow.backend.schemas import (
    GraphEdge,
    GraphNode,
    GraphPosition,
    PipelineCreateRequest,
    PipelineGraph,
)
from govoplan_dataflow.backend.service import (
    cancel_pipeline_run,
    create_pipeline,
    start_pipeline_run,
)


def _principal() -> ApiPrincipal:
    return ApiPrincipal(
        principal=PrincipalRef(
            account_id="account-1",
            membership_id="membership-1",
            tenant_id="tenant-1",
            scopes=frozenset({"dataflow:pipeline:run"}),
        ),
        account=object(),
        user=object(),
    )


def _graph() -> PipelineGraph:
    return PipelineGraph(
        nodes=[
            GraphNode(
                id="source",
                type="source.inline",
                label="Input",
                position=GraphPosition(x=0, y=0),
                config={
                    "source_name": "input_rows",
                    "rows": [{"id": 1}, {"id": 2}],
                },
            ),
            GraphNode(
                id="output",
                type="output",
                label="Output",
                position=GraphPosition(x=200, y=0),
                config={},
            ),
        ],
        edges=[
            GraphEdge(
                id="source-output",
                source="source",
                target="output",
            )
        ],
    )


class _AutomationProvider:
    def resolve_automation_principal(self, _session, *, request):
        return AutomationPrincipalResolution(
            allowed=True,
            principal=_principal(),
            granted_scopes=request.grant_scopes,
            provenance={"status": "rechecked"},
        )


class _Registry:
    def __init__(self) -> None:
        self.provider = _AutomationProvider()

    def has_capability(self, name: str) -> bool:
        return name == CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER

    def capability(self, name: str):
        if not self.has_capability(name):
            raise KeyError(name)
        return self.provider


class DataflowRunWorkerTests(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__,
                DataflowPipelineDeployment.__table__,
            ],
        )
        self.Session = sessionmaker(bind=self.engine)
        self.session: Session = self.Session()
        self.pipeline = create_pipeline(
            self.session,
            tenant_id="tenant-1",
            actor_id="account-1",
            payload=PipelineCreateRequest(
                name="Worker flow",
                status="active",
                graph=_graph(),
                editor_mode="graph",
            ),
        )
        self.session.commit()

    def tearDown(self) -> None:
        self.session.close()
        Base.metadata.drop_all(
            self.engine,
            tables=[
                DataflowPipelineDeployment.__table__,
                DataflowRun.__table__,
                DataflowPipelineRevision.__table__,
                DataflowPipeline.__table__,
            ],
        )
        self.engine.dispose()

    def _queue(self, key: str) -> DataflowRun:
        run, replayed = start_pipeline_run(
            self.session,
            tenant_id="tenant-1",
            actor_id="account-1",
            principal=_principal(),
            registry=_Registry(),
            request=DataflowRunRequest(
                pipeline_ref=f"pipeline:{self.pipeline.id}",
                revision=1,
                idempotency_key=key,
            ),
            defer_execution=True,
        )
        self.session.commit()
        self.assertFalse(replayed)
        return run

    def test_worker_claims_authorizes_and_executes_queued_run(self) -> None:
        run = self._queue("worker-success")
        self.assertEqual("queued", run.status)
        self.assertIsNone(run.started_at)

        result = dispatch_pending_runs(
            self.session,
            registry=_Registry(),
            worker_id="test-worker",
        )

        self.session.refresh(run)
        self.assertEqual(1, result["claimed"])
        self.assertEqual(1, result["succeeded"])
        self.assertEqual("succeeded", run.status)
        self.assertEqual(2, run.output_row_count)
        self.assertEqual(1, run.attempts)
        self.assertEqual(100, run.progress_percent)
        self.assertIsNone(run.worker_id)
        self.assertEqual(
            "rechecked",
            run.authorization_["last_resolution"]["status"],
        )

    def test_cancelled_queued_run_is_never_claimed(self) -> None:
        run = self._queue("worker-cancel")
        cancel_pipeline_run(
            self.session,
            tenant_id="tenant-1",
            run_ref=f"dataflow-run:{run.id}",
        )
        self.session.commit()

        result = dispatch_pending_runs(
            self.session,
            registry=_Registry(),
        )

        self.assertEqual(0, result["claimed"])
        self.session.refresh(run)
        self.assertEqual("cancelled", run.status)

    def test_expired_worker_lease_is_recovered(self) -> None:
        run = self._queue("worker-recover")
        run.status = "running"
        run.attempts = 1
        run.worker_id = "lost-worker"
        run.lease_expires_at = utcnow() - timedelta(minutes=1)
        self.session.commit()

        result = dispatch_pending_runs(
            self.session,
            registry=_Registry(),
        )

        self.session.refresh(run)
        self.assertEqual(1, result["recovered"])
        self.assertEqual("succeeded", run.status)
        self.assertEqual(2, run.attempts)

    def test_retention_purges_payload_but_keeps_run_evidence(self) -> None:
        run = self._queue("worker-retention")
        dispatch_pending_runs(self.session, registry=_Registry())
        run.retention_until = utcnow() - timedelta(seconds=1)
        self.session.commit()

        result = purge_expired_runs(self.session)

        self.session.refresh(run)
        self.assertEqual(1, result["purged"])
        self.assertEqual({}, run.request_)
        self.assertEqual("succeeded", run.status)
        self.assertEqual(2, run.output_row_count)
        self.assertIsNotNone(run.purged_at)


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