MockClient

@export_module("autogen.a2a")
def MockClient(  # noqa: N802
    response_message: str | dict[str, Any] | TextPart | DataPart | Part,
    *,
    extended_agent_card: AgentCard | None = None,
) -> HttpxClientFactory:
    """Create a mock HTTP client for testing A2A agent interactions.

    This function creates a mock HTTP client that simulates responses from an A2A agent server.
    It handles both agent card requests and message sending requests with configurable responses.

    Args:
        response_message: The message to return in response to SendMessage requests.
        extended_agent_card: Optional extended agent card to serve at the
            authenticated extended-card endpoint. When provided, the public
            card advertises ``supports_authenticated_extended_card=True`` so
            clients fetch the extended one.

    Returns:
        An HttpxClientFactory configured with a mock transport that handles requests
        to agent card endpoints and message sending endpoints.

    Example:
        >>> client = MockClient("Hello, world!")
        >>> agent = A2aRemoteAgent(name="remote", url="http://fake", client=client)
    """
    if isinstance(response_message, Part):
        parts = [response_message]
    elif isinstance(response_message, (DataPart, TextPart)):
        parts = [Part(root=response_message)]
    elif isinstance(response_message, str):
        parts = [Part(root=DataPart(data={"role": "assistant", "content": response_message}))]
    elif isinstance(response_message, dict):
        parts = [Part(root=DataPart(data={"role": "assistant", **response_message}))]
    else:
        raise ValueError(f"Invalid message type: {type(response_message)}")

    public_card = AgentCard(
        capabilities=AgentCapabilities(streaming=False),
        default_input_modes=["text"],
        default_output_modes=["text"],
        name="mock_agent",
        description="mock_agent",
        url="http://localhost:8000",
        supports_authenticated_extended_card=extended_agent_card is not None,
        version="0.1.0",
        skills=[],
    )

    public_card_paths = (
        AGENT_CARD_WELL_KNOWN_PATH,  # 1.0 well-known
        PREV_AGENT_CARD_WELL_KNOWN_PATH,  # v0.3 legacy well-known
    )
    extended_card_paths = (
        EXTENDED_AGENT_CARD_PATH,  # 1.0 extended card
        "/agent/authenticatedExtendedCard",  # v0.3 legacy extended
    )

    async def mock_handler(request: Request) -> Response:
        if request.url.path in public_card_paths:
            return Response(status_code=200, content=public_card.model_dump_json())

        if request.url.path in extended_card_paths:
            if extended_agent_card is None:
                return Response(status_code=404)
            return Response(status_code=200, content=extended_agent_card.model_dump_json())

        return Response(
            status_code=200,
            content=SendMessageSuccessResponse(
                result=Message(
                    message_id=str(uuid4()),
                    role=Role.agent,
                    parts=parts,
                ),
            ).model_dump_json(),
        )

    return HttpxClientFactory(transport=MockTransport(handler=mock_handler))