diff --git a/src/backend/tests/unit/build_utils.py b/src/backend/tests/unit/build_utils.py index abad2d76998a..88143d6d9a44 100644 --- a/src/backend/tests/unit/build_utils.py +++ b/src/backend/tests/unit/build_utils.py @@ -13,13 +13,25 @@ async def create_flow(client: AsyncClient, flow_data: str, headers: dict[str, st async def build_flow( - client: AsyncClient, flow_id: UUID, headers: dict[str, str], json: dict[str, Any] | None = None + client: AsyncClient, + flow_id: UUID, + headers: dict[str, str], + json: dict[str, Any] | None = None, + event_delivery: str | None = None, ) -> dict[str, Any]: - """Start a flow build and return the job_id.""" + """Start a flow build and return the job_id or events for direct delivery.""" if json is None: json = {} - response = await client.post(f"api/v1/build/{flow_id}/flow", json=json, headers=headers) + + # Add event_delivery to query params if specified + params = {} + if event_delivery: + params["event_delivery"] = event_delivery + + response = await client.post(f"api/v1/build/{flow_id}/flow", json=json, headers=headers, params=params) assert response.status_code == codes.OK + if event_delivery == "direct": + return response return response.json() diff --git a/src/backend/tests/unit/test_chat_endpoint.py b/src/backend/tests/unit/test_chat_endpoint.py index 9e97c76897f2..3fd99c09bb7a 100644 --- a/src/backend/tests/unit/test_chat_endpoint.py +++ b/src/backend/tests/unit/test_chat_endpoint.py @@ -358,3 +358,16 @@ async def mock_cancel_flow_build_with_cancelled_error(*_args, **_kwargs): finally: # Restore the original function to avoid affecting other tests monkeypatch.setattr(langflow.api.v1.chat, "cancel_flow_build", original_cancel_flow_build) + + +@pytest.mark.benchmark +async def test_build_flow_direct(client, json_memory_chatbot_no_llm, logged_in_headers): + """Test the build flow endpoint with direct event delivery.""" + # First create the flow + flow_id = await create_flow(client, json_memory_chatbot_no_llm, logged_in_headers) + + # Start the build with direct event delivery + response = await build_flow(client, flow_id, logged_in_headers, event_delivery="direct") + # Use the same consume_and_assert_stream function to verify the events + # We pass None as job_id since direct delivery doesn't use job_ids + await consume_and_assert_stream(response, None)