from __future__ import annotations import unittest from govoplan_core.core.dataflows import ( CAPABILITY_DATAFLOW_RUN_LIFECYCLE, CAPABILITY_DATAFLOW_RUN_WORKER, DataflowRunDescriptor, DataflowRunLifecycleProvider, DataflowRunWorker, dataflow_run_lifecycle, dataflow_run_worker, ) from govoplan_core.core.modules import ModuleContext, ModuleManifest from govoplan_core.core.registry import PlatformRegistry class _Provider: def start_run(self, session, principal, *, request): del session, principal return DataflowRunDescriptor( ref="dataflow-run:1", pipeline_ref=request.pipeline_ref, revision=request.revision, status="succeeded", definition_hash="abc", executor_version="test", ) def get_run(self, session, principal, *, run_ref): del session, principal if run_ref != "dataflow-run:1": return None return DataflowRunDescriptor( ref=run_ref, pipeline_ref="pipeline:1", revision=1, status="succeeded", definition_hash="abc", executor_version="test", ) def cancel_run(self, session, principal, *, run_ref): del session, principal return DataflowRunDescriptor( ref=run_ref, pipeline_ref="pipeline:1", revision=1, status="cancelled", definition_hash="abc", executor_version="test", ) class _Worker: def dispatch_pending( self, session, *, now=None, limit=10, worker_id=None, ): return {"claimed": 0} def purge_expired(self, session, *, now=None, limit=500): return {"purged": 0} class DataflowContractTests(unittest.TestCase): def test_run_lifecycle_is_runtime_checkable_and_resolved(self) -> None: provider = _Provider() worker = _Worker() self.assertIsInstance(provider, DataflowRunLifecycleProvider) self.assertIsInstance(worker, DataflowRunWorker) registry = PlatformRegistry() registry.register( ModuleManifest( id="dataflow_contract_test", name="Dataflow contract test", version="test", capability_factories={ CAPABILITY_DATAFLOW_RUN_LIFECYCLE: lambda context: provider, CAPABILITY_DATAFLOW_RUN_WORKER: lambda context: worker, }, ) ) registry.configure_capability_context( ModuleContext(registry=registry, settings=object()) ) self.assertIs(provider, dataflow_run_lifecycle(registry)) self.assertIs(worker, dataflow_run_worker(registry)) self.assertIsNone(dataflow_run_lifecycle(PlatformRegistry())) self.assertIsNone(dataflow_run_worker(PlatformRegistry())) if __name__ == "__main__": unittest.main()