from __future__ import annotations from dataclasses import replace 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 ( CAPABILITY_DATAFLOW_RUN_LIFECYCLE, DataflowRunDescriptor, ) from govoplan_core.db.base import Base from govoplan_workflow.backend.db.models import ( WorkflowDefinition, WorkflowDefinitionRevision, WorkflowInstance, WorkflowInstanceEvent, WorkflowInstanceStep, ) from govoplan_workflow.backend.instance_service import ( SqlWorkflowRuntimeWorker, cancel_instance, instance_response, reconcile_instance, resolve_step, start_instance, ) from govoplan_workflow.backend.schemas import ( WorkflowDefinitionCreateRequest, WorkflowEdge, WorkflowGraph, WorkflowInstanceStartRequest, WorkflowNode, WorkflowStepActionRequest, ) from govoplan_workflow.backend.service import ( WorkflowConflictError, activate_definition, create_definition, ) def principal() -> ApiPrincipal: return ApiPrincipal( principal=PrincipalRef( account_id="account-1", membership_id="membership-1", tenant_id="tenant-1", scopes=frozenset( { "workflow:definition:read", "workflow:instance:read", "workflow:instance:start", "workflow:instance:transition", "dataflow:pipeline:run", } ), ), account=object(), user=object(), ) def runtime_graph() -> WorkflowGraph: return WorkflowGraph( nodes=[ WorkflowNode( id="start", type="workflow.start.manual", label="Start", config={"input_schema_ref": ""}, ), WorkflowNode( id="flow", type="workflow.dataflow", label="Prepare evidence", config={ "pipeline_ref": "pipeline:pipeline-1", "revision": 3, "environment": "development", "row_limit": 250, "publication_target_ref": "", "warning_policy": "review", "input_mapping": {}, }, ), WorkflowNode( id="complete", type="workflow.end.completed", label="Complete", config={"output_mapping": {}}, ), WorkflowNode( id="cancelled", type="workflow.end.cancelled", label="Rejected", config={"reason": "Rejected during review"}, ), ], edges=[ WorkflowEdge( id="start-flow", source="start", target="flow", ), WorkflowEdge( id="flow-complete", source="flow", source_port="success", target="complete", ), WorkflowEdge( id="flow-warning", source="flow", source_port="warning", target="complete", ), WorkflowEdge( id="flow-review", source="flow", source_port="review_required", target="complete", ), WorkflowEdge( id="flow-failure", source="flow", source_port="failure", target="cancelled", ), ], ) class FakeDataflowLifecycle: def __init__(self) -> None: self.runs: dict[str, DataflowRunDescriptor] = {} self.requests = [] self.cancelled: list[str] = [] def start_run(self, _session, _principal, *, request): self.requests.append(request) run_ref = f"run:{len(self.requests)}" descriptor = DataflowRunDescriptor( ref=run_ref, pipeline_ref=request.pipeline_ref, revision=request.revision, status="queued", definition_hash="definition-hash", executor_version="test", metadata={"progress_percent": 0, "progress_phase": "queued"}, ) self.runs[run_ref] = descriptor return descriptor def get_run(self, _session, _principal, *, run_ref): return self.runs.get(run_ref) def cancel_run(self, _session, _principal, *, run_ref): descriptor = self.runs[run_ref] descriptor = replace(descriptor, status="cancelled") self.runs[run_ref] = descriptor self.cancelled.append(run_ref) return descriptor def finish( self, run_ref: str, *, diagnostics: list[dict[str, object]] | None = None, ) -> None: self.runs[run_ref] = replace( self.runs[run_ref], status="succeeded", output_publication_ref="publication:1", output_datasource_ref="datasource:1", output_materialization_ref="materialization:1", input_row_count=12, output_row_count=10, metadata={"diagnostics": diagnostics or []}, ) def fail(self, run_ref: str) -> None: self.runs[run_ref] = replace( self.runs[run_ref], status="failed", error="Data quality gate failed.", ) class FakeAutomationProvider: def __init__(self) -> None: self.requests = [] def resolve_automation_principal(self, _session, *, request): self.requests.append(request) return AutomationPrincipalResolution( allowed=True, principal=principal(), granted_scopes=request.grant_scopes, provenance={"status": "rechecked"}, ) class Registry: def __init__( self, dataflow: FakeDataflowLifecycle, automation: FakeAutomationProvider | None = None, ) -> None: self.dataflow = dataflow self.automation = automation def has_capability(self, name: str) -> bool: return name == CAPABILITY_DATAFLOW_RUN_LIFECYCLE or ( name == CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER and self.automation is not None ) def capability(self, name: str): if name == CAPABILITY_DATAFLOW_RUN_LIFECYCLE: return self.dataflow if ( name == CAPABILITY_AUTH_AUTOMATION_PRINCIPAL_PROVIDER and self.automation is not None ): return self.automation raise KeyError(name) class WorkflowInstanceServiceTests(unittest.TestCase): def setUp(self) -> None: self.engine = create_engine("sqlite:///:memory:") Base.metadata.create_all( self.engine, tables=[ WorkflowDefinition.__table__, WorkflowDefinitionRevision.__table__, WorkflowInstance.__table__, WorkflowInstanceStep.__table__, WorkflowInstanceEvent.__table__, ], ) self.Session = sessionmaker(bind=self.engine) self.session: Session = self.Session() self.dataflow = FakeDataflowLifecycle() self.registry = Registry(self.dataflow) self.definition = create_definition( self.session, tenant_id="tenant-1", actor_id="account-1", payload=WorkflowDefinitionCreateRequest( name="Monthly governed processing", graph=runtime_graph(), ), ) activate_definition( self.session, tenant_id="tenant-1", definition_id=self.definition.id, actor_id="account-1", revision=1, ) self.session.commit() def tearDown(self) -> None: self.session.close() Base.metadata.drop_all( self.engine, tables=[ WorkflowInstanceEvent.__table__, WorkflowInstanceStep.__table__, WorkflowInstance.__table__, WorkflowDefinitionRevision.__table__, WorkflowDefinition.__table__, ], ) self.engine.dispose() def _start(self, key: str = "request-1") -> WorkflowInstance: instance, replayed = start_instance( self.session, tenant_id="tenant-1", definition_id=self.definition.id, actor_id="account-1", principal=principal(), registry=self.registry, payload=WorkflowInstanceStartRequest( idempotency_key=key, input={"case_id": "case-1"}, correlation_id="correlation-1", ), ) self.assertFalse(replayed) return instance def test_start_pins_revision_and_replays_idempotently(self) -> None: instance = self._start() replayed, was_replayed = start_instance( self.session, tenant_id="tenant-1", definition_id=self.definition.id, actor_id="account-1", principal=principal(), registry=self.registry, payload=WorkflowInstanceStartRequest( idempotency_key="request-1", input={"case_id": "case-1"}, correlation_id="correlation-1", ), ) self.assertTrue(was_replayed) self.assertEqual(instance.id, replayed.id) self.assertEqual("waiting", instance.status) self.assertEqual(1, len(self.dataflow.requests)) response = instance_response(self.session, instance) self.assertEqual([1, 2], [step.sequence for step in response.steps]) self.assertEqual("run:1", response.steps[-1].external_ref) self.assertGreaterEqual(len(response.events), 4) with self.assertRaises(WorkflowConflictError): start_instance( self.session, tenant_id="tenant-1", definition_id=self.definition.id, actor_id="account-1", principal=principal(), registry=self.registry, payload=WorkflowInstanceStartRequest( idempotency_key="request-1", input={"case_id": "another-case"}, correlation_id="correlation-1", ), ) def test_reconcile_completes_with_stable_dataflow_output_refs(self) -> None: instance = self._start() self.dataflow.finish("run:1") changed = reconcile_instance( self.session, instance=instance, principal=principal(), registry=self.registry, actor_id="account-1", ) response = instance_response(self.session, instance) self.assertTrue(changed) self.assertEqual("completed", response.status) flow_output = response.context["steps"]["flow"] self.assertEqual("publication:1", flow_output["output_publication_ref"]) self.assertEqual("datasource:1", flow_output["output_datasource_ref"]) self.assertEqual( "materialization:1", flow_output["output_materialization_ref"], ) self.assertEqual("workflow.instance.completed", response.events[-1].kind) def test_warning_requires_review_and_approve_resumes(self) -> None: instance = self._start() self.dataflow.finish( "run:1", diagnostics=[ { "severity": "warning", "code": "review.required", "message": "Verify unmatched records.", } ], ) reconcile_instance( self.session, instance=instance, principal=principal(), registry=self.registry, ) step_id = str(instance.current_step_id) self.assertEqual( "review_required", instance_response(self.session, instance).steps[-1].handoff["state"], ) resolved = resolve_step( self.session, tenant_id="tenant-1", instance_id=instance.id, step_id=step_id, actor_id="account-1", principal=principal(), registry=self.registry, payload=WorkflowStepActionRequest( action="approve", comment="Evidence verified.", evidence=["publication:1"], ), ) self.assertEqual("completed", resolved.status) def test_failure_can_retry_and_reject_invalid_actions(self) -> None: instance = self._start() self.dataflow.fail("run:1") reconcile_instance( self.session, instance=instance, principal=principal(), registry=self.registry, ) step_id = str(instance.current_step_id) with self.assertRaises(WorkflowConflictError): resolve_step( self.session, tenant_id="tenant-1", instance_id=instance.id, step_id=step_id, actor_id="account-1", principal=principal(), registry=self.registry, payload=WorkflowStepActionRequest(action="approve"), ) retried = resolve_step( self.session, tenant_id="tenant-1", instance_id=instance.id, step_id=step_id, actor_id="account-1", principal=principal(), registry=self.registry, payload=WorkflowStepActionRequest(action="retry"), ) self.assertEqual("waiting", retried.status) self.assertEqual(2, len(self.dataflow.requests)) self.assertEqual("run:2", retried.steps[-1].external_ref) self.assertEqual("superseded", retried.steps[-2].status) def test_cancel_propagates_to_linked_dataflow(self) -> None: instance = self._start() cancelled = cancel_instance( self.session, tenant_id="tenant-1", instance_id=instance.id, actor_id="account-1", principal=principal(), registry=self.registry, ) self.assertEqual("cancelled", cancelled.status) self.assertEqual(["run:1"], self.dataflow.cancelled) def test_worker_rechecks_authorization_before_reconciling(self) -> None: instance = self._start() self.session.commit() self.dataflow.finish("run:1") automation = FakeAutomationProvider() worker = SqlWorkflowRuntimeWorker( registry=Registry(self.dataflow, automation), ) summary = worker.reconcile_pending(self.session) self.assertEqual(1, summary["advanced"]) self.assertEqual("completed", instance.status) self.assertEqual(1, len(automation.requests)) self.assertEqual( "rechecked", instance.authorization_["last_resolution"]["status"], ) if __name__ == "__main__": unittest.main()