Skip prefect when using BusSimulator

Author: dgarros

backend/infrahub/message_bus/operations/__init__.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 import ujson
+from prefect import Flow
 
 from infrahub.message_bus import RPCErrorResponse, messages
 from infrahub.message_bus.operations import (
@@ -81,12 +82,17 @@
 }
 
 
-async def execute_message(routing_key: str, message_body: bytes, service: InfrahubServices) -> Optional[MessageTTL]:
+async def execute_message(
+    routing_key: str, message_body: bytes, service: InfrahubServices, skip_flow: bool = False
+) -> Optional[MessageTTL]:
     message_data = ujson.loads(message_body)
     message = messages.MESSAGE_MAP[routing_key](**message_data)
     message.set_log_data(routing_key=routing_key)
     try:
-        await COMMAND_MAP[routing_key](message=message, service=service)
+        func = COMMAND_MAP[routing_key]
+        if skip_flow and isinstance(func, Flow):
+            func = func.fn
+        await func(message=message, service=service)
     except Exception as exc:  # pylint: disable=broad-except
         if message.reply_requested:
             response = RPCErrorResponse(errors=[str(exc)], initial_message=message.model_dump())

backend/tests/adapters/message_bus.py

@@ -50,7 +50,7 @@ async def publish(
         if routing_key not in self.messages_per_routing_key:
             self.messages_per_routing_key[routing_key] = []
         self.messages_per_routing_key[routing_key].append(message)
-        await execute_message(routing_key=routing_key, message_body=message.body, service=self.service)
+        await execute_message(routing_key=routing_key, message_body=message.body, service=self.service, skip_flow=True)
 
     async def reply(self, message: InfrahubMessage, routing_key: str) -> None:
         correlation_id = message.meta.correlation_id or "default"