dify/api/tests/unit_tests/clients/agent_backend/test_client.py

from collections.abc import Callable, Iterator
from typing import override

import pytest
from dify_agent.client import DifyAgentHTTPError, DifyAgentStreamError, DifyAgentTimeoutError, DifyAgentValidationError
from dify_agent.layers.execution_context import DifyExecutionContextLayerConfig
from dify_agent.protocol import (
    CancelRunRequest,
    CancelRunResponse,
    CreateRunRequest,
    CreateRunResponse,
    RunEvent,
    RunStartedEvent,
    RunStatusResponse,
)

from clients.agent_backend import (
    AgentBackendHTTPError,
    AgentBackendModelConfig,
    AgentBackendRunRequestBuilder,
    AgentBackendStreamError,
    AgentBackendTransportError,
    AgentBackendValidationError,
    AgentBackendWorkflowNodeRunInput,
    DifyAgentBackendRunClient,
)


def _request():
    return AgentBackendRunRequestBuilder().build_for_workflow_node(
        AgentBackendWorkflowNodeRunInput(
            model=AgentBackendModelConfig(
                plugin_id="langgenius/openai",
                model_provider="openai",
                model="gpt-test",
            ),
            execution_context=DifyExecutionContextLayerConfig(
                tenant_id="tenant-1",
                user_from="account",
                agent_mode="workflow_run",
                invoke_from="debugger",
            ),
            workflow_node_job_prompt="Do the task.",
            user_prompt="hello",
        )
    )


class _SuccessfulClient:
    stream_options: tuple[int | None, float | None, Callable[[], bool] | None] | None = None

    def create_run_sync(self, request: CreateRunRequest) -> CreateRunResponse:
        assert isinstance(request, CreateRunRequest)
        return CreateRunResponse(run_id="run-1", status="running")

    def cancel_run_sync(self, run_id: str, request: CancelRunRequest | None = None) -> CancelRunResponse:
        del request
        return CancelRunResponse(run_id=run_id, status="cancelled")

    def stream_events_sync(
        self,
        run_id: str,
        *,
        after: str | None = None,
        max_reconnects: int | None = None,
        timeout_seconds: float | None = None,
        should_stop: Callable[[], bool] | None = None,
    ) -> Iterator[RunEvent]:
        del after
        self.stream_options = (max_reconnects, timeout_seconds, should_stop)
        yield RunStartedEvent(id="1-0", run_id=run_id)

    def wait_run_sync(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse:
        del timeout_seconds
        return RunStatusResponse.model_validate(
            {
                "run_id": run_id,
                "status": "succeeded",
                "created_at": "2026-01-01T00:00:00+00:00",
                "updated_at": "2026-01-01T00:00:00+00:00",
            }
        )


def test_dify_agent_backend_run_client_delegates_sync_methods():
    wrapped = _SuccessfulClient()
    client = DifyAgentBackendRunClient(wrapped, stream_max_reconnects=2, stream_timeout_seconds=45)

    def should_stop() -> bool:
        return False

    created = client.create_run(_request())
    cancelled = client.cancel_run(created.run_id)
    events = list(client.stream_events(created.run_id, should_stop=should_stop))
    status = client.wait_run(created.run_id)

    assert created.run_id == "run-1"
    assert cancelled.status == "cancelled"
    assert events[0].type == "run_started"
    assert status.status == "succeeded"
    assert wrapped.stream_options == (2, 45, should_stop)


def test_dify_agent_backend_run_client_maps_validation_error():
    class InvalidClient(_SuccessfulClient):
        @override
        def create_run_sync(self, request: CreateRunRequest) -> CreateRunResponse:
            raise DifyAgentValidationError(detail={"field": "bad"})

    with pytest.raises(AgentBackendValidationError) as exc_info:
        DifyAgentBackendRunClient(InvalidClient()).create_run(_request())

    assert exc_info.value.detail == {"field": "bad"}


def test_dify_agent_backend_run_client_maps_http_error():
    class HTTPErrorClient(_SuccessfulClient):
        @override
        def create_run_sync(self, request: CreateRunRequest) -> CreateRunResponse:
            raise DifyAgentHTTPError(status_code=503, detail="unavailable")

    with pytest.raises(AgentBackendHTTPError) as exc_info:
        DifyAgentBackendRunClient(HTTPErrorClient()).create_run(_request())

    assert exc_info.value.status_code == 503
    assert exc_info.value.detail == "unavailable"


def test_dify_agent_backend_run_client_maps_timeout_error():
    class TimeoutClient(_SuccessfulClient):
        @override
        def wait_run_sync(self, run_id: str, *, timeout_seconds: float | None = None) -> RunStatusResponse:
            raise DifyAgentTimeoutError("timeout")

    with pytest.raises(AgentBackendTransportError) as exc_info:
        DifyAgentBackendRunClient(TimeoutClient()).wait_run("run-1")

    assert str(exc_info.value) == "timeout"


def test_dify_agent_backend_run_client_maps_stream_error():
    class StreamClient(_SuccessfulClient):
        @override
        def stream_events_sync(
            self,
            run_id: str,
            *,
            after: str | None = None,
            max_reconnects: int | None = None,
            timeout_seconds: float | None = None,
            should_stop: Callable[[], bool] | None = None,
        ) -> Iterator[RunEvent]:
            del run_id, after, max_reconnects, timeout_seconds, should_stop
            raise DifyAgentStreamError("bad stream")
            yield

    with pytest.raises(AgentBackendStreamError) as exc_info:
        list(DifyAgentBackendRunClient(StreamClient()).stream_events("run-1"))

    assert str(exc_info.value) == "bad stream"