Add integration tests for Prefect worker

Author: dgarros

backend/infrahub/api/schema.py

@@ -279,11 +279,13 @@ async def load_schema(
             schema_branch=candidate_schema,
             constraints=result.constraints,
         )
-        error_messages = await service.workflow.execute(  # type: ignore[var-annotated]
-            workflow=SCHEMA_VALIDATE_MIGRATION, message=validate_migration_data
+        error_messages = await service.workflow.execute_workflow(
+            workflow=SCHEMA_VALIDATE_MIGRATION,
+            expected_return=list[str],
+            parameters={"message": validate_migration_data},
         )
-        if error_messages:  # type: ignore[has-type]
-            raise SchemaNotValidError(message=",\n".join(error_messages))  # type: ignore[has-type]
+        if error_messages:
+            raise SchemaNotValidError(message=",\n".join(error_messages))
 
         # ----------------------------------------------------------
         # Update the schema
@@ -320,8 +322,8 @@ async def load_schema(
             previous_schema=origin_schema,
             migrations=result.migrations,
         )
-        migration_error_msgs = await service.workflow.execute(  # type: ignore[var-annotated]
-            workflow=SCHEMA_APPLY_MIGRATION, message=apply_migration_data
+        migration_error_msgs = await service.workflow.execute_workflow(
+            workflow=SCHEMA_APPLY_MIGRATION, expected_return=list[str], parameters={"message": apply_migration_data}
         )
 
         if migration_error_msgs:
@@ -368,8 +370,8 @@ async def check_schema(
         schema_branch=candidate_schema,
         constraints=result.constraints,
     )
-    error_messages = await service.workflow.execute(  # type: ignore[var-annotated]
-        workflow=SCHEMA_VALIDATE_MIGRATION, message=validate_migration_data
+    error_messages = await service.workflow.execute_workflow(
+        workflow=SCHEMA_VALIDATE_MIGRATION, expected_return=list[str], parameters={"message": validate_migration_data}
     )
     if error_messages:
         raise SchemaNotValidError(message=",\n".join(error_messages))

backend/infrahub/api/transformation.py

@@ -146,6 +146,7 @@ async def transform_jinja2(
 
     service: InfrahubServices = request.app.state.service
 
-    response: str = await service.workflow.execute(workflow=TRANSFORM_JINJA2_RENDER, message=message)  # type: ignore[arg-type]
-
+    response = await service.workflow.execute_workflow(
+        workflow=TRANSFORM_JINJA2_RENDER, expected_return=str, parameters={"message": message}
+    )
     return PlainTextResponse(content=response)

backend/infrahub/cli/db.py

@@ -243,8 +243,10 @@ async def update_core_schema(  # pylint: disable=too-many-statements
                 schema_branch=candidate_schema,
                 constraints=result.constraints,
             )
-            error_messages = await service.workflow.execute(  # type: ignore[var-annotated]
-                workflow=SCHEMA_VALIDATE_MIGRATION, message=validate_migration_data
+            error_messages = await service.workflow.execute_workflow(
+                workflow=SCHEMA_VALIDATE_MIGRATION,
+                expected_return=list[str],
+                parameters={"message": validate_migration_data},
             )
             if error_messages:
                 rprint(f"{error_badge} | Unable to update the schema, due to failed validations")
@@ -286,8 +288,8 @@ async def update_core_schema(  # pylint: disable=too-many-statements
                 previous_schema=origin_schema,
                 migrations=result.migrations,
             )
-            migration_error_msgs = await service.workflow.execute(  # type: ignore[var-annotated]
-                workflow=SCHEMA_APPLY_MIGRATION, message=apply_migration_data
+            migration_error_msgs = await service.workflow.execute_workflow(
+                workflow=SCHEMA_APPLY_MIGRATION, expected_return=list[str], parameters={"message": apply_migration_data}
             )
 
             if migration_error_msgs:

backend/infrahub/cli/tasks.py

@@ -46,5 +46,7 @@ async def execute(
         worker = WorkflowWorkerExecution()
         await DUMMY_FLOW.save(client=client, work_pool=WorkerPoolDefinition(name="testing", worker_type="process"))
 
-        result = await worker.execute(workflow=DUMMY_FLOW, data=DummyInput(firstname="John", lastname="Doe"))  # type: ignore[var-annotated]
+        result = await worker.execute_workflow(
+            workflow=DUMMY_FLOW, parameters={"data": DummyInput(firstname="John", lastname="Doe")}
+        )  # type: ignore[var-annotated]
         print(result)

backend/infrahub/message_bus/operations/refresh/webhook.py

@@ -11,6 +11,10 @@ async def configuration(
     message: messages.RefreshWebhookConfiguration,  # pylint: disable=unused-argument
     service: InfrahubServices,
 ) -> None:
+    if not service._client:
+        service.log.error("Client hasn't been initialized, can't refresh webhook")
+        return
+
     service.log.debug("Refreshing webhook configuration")
     standard_webhooks = await service.client.all(kind=InfrahubKind.STANDARDWEBHOOK)
     custom_webhooks = await service.client.all(kind=InfrahubKind.CUSTOMWEBHOOK)

backend/infrahub/services/adapters/workflow/__init__.py

@@ -1,10 +1,10 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, ParamSpec, TypeVar
+from typing import TYPE_CHECKING, Any, Callable, ParamSpec, TypeVar, overload
 
 if TYPE_CHECKING:
     from infrahub.services import InfrahubServices
-    from infrahub.workflows.models import WorkflowDefinition
+    from infrahub.workflows.models import WorkflowDefinition, WorkflowInfo
 
 Return = TypeVar("Return")
 Params = ParamSpec("Params")
@@ -16,10 +16,33 @@ class InfrahubWorkflow:
     async def initialize(self, service: InfrahubServices) -> None:
         """Initialize the Workflow engine"""
 
-    async def execute(
+    @overload
+    async def execute_workflow(
         self,
-        workflow: WorkflowDefinition | None = None,
-        function: Callable[..., Awaitable[Return]] | None = None,
-        **kwargs: Any,
-    ) -> Return:
+        workflow: WorkflowDefinition,
+        expected_return: type[Return],
+        parameters: dict[str, Any] | None = ...,
+    ) -> Return: ...
+
+    @overload
+    async def execute_workflow(
+        self,
+        workflow: WorkflowDefinition,
+        expected_return: None = ...,
+        parameters: dict[str, Any] | None = ...,
+    ) -> Any: ...
+
+    async def execute_workflow(
+        self,
+        workflow: WorkflowDefinition,
+        expected_return: type[Return] | None = None,
+        parameters: dict[str, Any] | None = None,
+    ) -> Any:
+        raise NotImplementedError()
+
+    async def submit_workflow(
+        self,
+        workflow: WorkflowDefinition,
+        parameters: dict[str, Any] | None = None,
+    ) -> WorkflowInfo:
         raise NotImplementedError()

backend/infrahub/services/adapters/workflow/local.py

@@ -1,20 +1,25 @@
-from typing import Any, Awaitable, Callable
+import uuid
+from typing import Any
 
-from infrahub.workflows.models import WorkflowDefinition
+from infrahub.workflows.models import WorkflowDefinition, WorkflowInfo
 
 from . import InfrahubWorkflow, Return
 
 
 class WorkflowLocalExecution(InfrahubWorkflow):
-    async def execute(
+    async def execute_workflow(
         self,
-        workflow: WorkflowDefinition | None = None,
-        function: Callable[..., Awaitable[Return]] | None = None,
-        **kwargs: Any,
-    ) -> Return:
-        if workflow:
-            fn = workflow.get_function()
-            return await fn(**kwargs)
-        if function:
-            return await function(**kwargs)
-        raise ValueError("either a workflow definition or a flow must be provided")
+        workflow: WorkflowDefinition,
+        expected_return: type[Return] | None = None,
+        parameters: dict[str, Any] | None = None,
+    ) -> Any:
+        fn = workflow.get_function()
+        return await fn(**parameters or {})
+
+    async def submit_workflow(
+        self,
+        workflow: WorkflowDefinition,
+        parameters: dict[str, Any] | None = None,
+    ) -> WorkflowInfo:
+        workflow.get_function()
+        return WorkflowInfo(id=uuid.uuid4())

backend/infrahub/services/adapters/workflow/worker.py

@@ -1,11 +1,12 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Awaitable, Callable
+from typing import TYPE_CHECKING, Any, overload
 
 from prefect.client.schemas import StateType
 from prefect.deployments import run_deployment
 
 from infrahub.workflows.initialization import setup_task_manager
+from infrahub.workflows.models import WorkflowInfo
 
 from . import InfrahubWorkflow, Return
 
@@ -23,23 +24,41 @@ async def initialize(self, service: InfrahubServices) -> None:
         if await service.component.is_primary_api():
             await setup_task_manager()
 
-    async def execute(
+    @overload
+    async def execute_workflow(
         self,
-        workflow: WorkflowDefinition | None = None,
-        function: Callable[..., Awaitable[Return]] | None = None,
-        **kwargs: Any,
-    ) -> Return:
-        if workflow:
-            response: FlowRun = await run_deployment(name=workflow.full_name, parameters=kwargs or {})  # type: ignore[return-value, misc]
-            if not response.state:
-                raise RuntimeError("Unable to read state from the response")
+        workflow: WorkflowDefinition,
+        expected_return: type[Return],
+        parameters: dict[str, Any] | None = ...,
+    ) -> Return: ...
 
-            if response.state.type == StateType.CRASHED:
-                raise RuntimeError(response.state.message)
+    @overload
+    async def execute_workflow(
+        self,
+        workflow: WorkflowDefinition,
+        expected_return: None = ...,
+        parameters: dict[str, Any] | None = ...,
+    ) -> Any: ...
+
+    async def execute_workflow(
+        self,
+        workflow: WorkflowDefinition,
+        expected_return: type[Return] | None = None,
+        parameters: dict[str, Any] | None = None,
+    ) -> Any:
+        response: FlowRun = await run_deployment(name=workflow.full_name, poll_interval=1, parameters=parameters or {})  # type: ignore[return-value, misc]
+        if not response.state:
+            raise RuntimeError("Unable to read state from the response")
 
-            return await response.state.result(raise_on_failure=True, fetch=True)  # type: ignore[call-overload]
+        if response.state.type == StateType.CRASHED:
+            raise RuntimeError(response.state.message)
 
-        if function:
-            return await function(**kwargs)
+        return await response.state.result(raise_on_failure=True, fetch=True)  # type: ignore[call-overload]
 
-        raise ValueError("either a workflow definition or a flow must be provided")
+    async def submit_workflow(
+        self,
+        workflow: WorkflowDefinition,
+        parameters: dict[str, Any] | None = None,
+    ) -> WorkflowInfo:
+        flow_run = await run_deployment(name=workflow.full_name, timeout=0, parameters=parameters or {})  # type: ignore[return-value, misc]
+        return WorkflowInfo.from_flow(flow_run=flow_run)

backend/infrahub/tasks/dummy.py

@@ -1,8 +1,20 @@
+from __future__ import annotations
+
 from prefect import flow, task
 from pydantic import BaseModel
 
 from infrahub.workflows.models import WorkflowDefinition
 
+
+class DummyInput(BaseModel):
+    firstname: str
+    lastname: str
+
+
+class DummyOutput(BaseModel):
+    full_name: str
+
+
 DUMMY_FLOW = WorkflowDefinition(
     name="dummy_flow",
     module="infrahub.tasks.dummy",
@@ -16,15 +28,6 @@
 )
 
 
-class DummyInput(BaseModel):
-    firstname: str
-    lastname: str
-
-
-class DummyOutput(BaseModel):
-    full_name: str
-
-
 @task
 async def aggregate_name(firstname: str, lastname: str) -> str:
     return f"{firstname}, {lastname}"