chore: import upstream snapshot with attribution

python/packages/ag-ui/tests/ag_ui/conftest.py

@@ -0,0 +1,350 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Shared test fixtures and stubs for AG-UI tests."""
+
+import sys
+from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable, Mapping, MutableSequence, Sequence
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Generic, Literal, TypedDict, cast, overload  # noqa: F401
+
+import pytest
+from agent_framework import (
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentSession,
+    BaseChatClient,
+    ChatOptions,
+    ChatResponse,
+    ChatResponseUpdate,
+    Content,
+    Message,
+    ServiceSessionId,
+    SupportsAgentRun,
+    SupportsChatGetResponse,
+)
+from agent_framework._clients import OptionsCoT
+from agent_framework._middleware import ChatMiddlewareLayer
+from agent_framework._tools import FunctionInvocationLayer
+from agent_framework._types import ResponseStream
+from agent_framework.observability import ChatTelemetryLayer
+
+if sys.version_info >= (3, 12):
+    from typing import override  # type: ignore # pragma: no cover
+else:
+    from typing_extensions import override  # type: ignore[import] # pragma: no cover
+
+StreamFn = Callable[..., AsyncIterable[ChatResponseUpdate]]
+ResponseFn = Callable[..., Awaitable[ChatResponse]]
+
+
+def pytest_configure() -> None:
+    """Ensure this test directory is on sys.path so helper modules can be imported by name."""
+    test_dir = str(Path(__file__).resolve().parent)
+    if test_dir not in sys.path:
+        sys.path.insert(0, test_dir)
+
+
+class StreamingChatClientStub(
+    FunctionInvocationLayer[OptionsCoT],
+    ChatMiddlewareLayer[OptionsCoT],
+    ChatTelemetryLayer[OptionsCoT],
+    BaseChatClient[OptionsCoT],
+    Generic[OptionsCoT],
+):
+    """Typed streaming stub that satisfies SupportsChatGetResponse."""
+
+    def __init__(self, stream_fn: StreamFn, response_fn: ResponseFn | None = None) -> None:
+        super().__init__(middleware=[])
+        self._stream_fn = stream_fn
+        self._response_fn = response_fn
+        self.last_session: AgentSession | None = None
+        self.last_service_session_id: str | ServiceSessionId | None = None
+
+    @overload
+    def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: Literal[False] = ...,
+        options: ChatOptions[Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse[Any]]: ...
+
+    @overload
+    def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: Literal[False] = ...,
+        options: OptionsCoT | ChatOptions[None] | None = ...,
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse[Any]]: ...
+
+    @overload
+    def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: Literal[True],
+        options: OptionsCoT | ChatOptions[Any] | None = ...,
+        **kwargs: Any,
+    ) -> ResponseStream[ChatResponseUpdate, ChatResponse[Any]]: ...
+
+    def get_response(
+        self,
+        messages: Sequence[Message],
+        *,
+        stream: bool = False,
+        options: OptionsCoT | ChatOptions[Any] | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
+        client_kwargs = kwargs.get("client_kwargs")
+        if isinstance(client_kwargs, Mapping):
+            self.last_session = cast(AgentSession | None, client_kwargs.get("session"))
+        else:
+            self.last_session = None
+        self.last_service_session_id = self.last_session.service_session_id if self.last_session else None
+        if stream:
+            return super().get_response(
+                messages=messages,
+                stream=True,
+                options=options,
+                **kwargs,
+            )
+        return super().get_response(
+            messages=messages,
+            stream=False,
+            options=options,
+            **kwargs,
+        )
+
+    @override
+    def _inner_get_response(
+        self,
+        *,
+        messages: Sequence[Message],
+        stream: bool = False,
+        options: Mapping[str, Any],
+        **kwargs: Any,
+    ) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
+        if stream:
+
+            def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
+                return ChatResponse.from_updates(updates)
+
+            return ResponseStream(self._stream_fn(messages, options, **kwargs), finalizer=_finalize)
+
+        return self._get_response_impl(messages, options, **kwargs)
+
+    async def _get_response_impl(
+        self, messages: Sequence[Message], options: Mapping[str, Any], **kwargs: Any
+    ) -> ChatResponse:
+        """Non-streaming implementation."""
+        if self._response_fn is not None:
+            return await self._response_fn(messages, options, **kwargs)
+
+        contents: list[Any] = []
+        async for update in self._stream_fn(list(messages), dict(options), **kwargs):
+            contents.extend(update.contents)
+
+        return ChatResponse(
+            messages=[Message(role="assistant", contents=contents)],
+            response_id="stub-response",
+        )
+
+
+def stream_from_updates(updates: list[ChatResponseUpdate]) -> StreamFn:
+    """Create a stream function that yields from a static list of updates."""
+
+    async def _stream(
+        messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
+    ) -> AsyncIterator[ChatResponseUpdate]:
+        for update in updates:
+            yield update
+
+    return _stream
+
+
+class StubAgent(SupportsAgentRun):
+    """Minimal SupportsAgentRun stub for orchestrator tests."""
+
+    def __init__(
+        self,
+        updates: list[AgentResponseUpdate] | None = None,
+        *,
+        agent_id: str = "stub-agent",
+        agent_name: str | None = "stub-agent",
+        default_options: Any | None = None,
+        client: Any | None = None,
+    ) -> None:
+        self.id = agent_id
+        self.name = agent_name
+        self.description = "stub agent"
+        self.updates = updates or [AgentResponseUpdate(contents=[Content.from_text(text="response")], role="assistant")]
+        self.default_options: dict[str, Any] = (
+            default_options if isinstance(default_options, dict) else {"tools": None, "response_format": None}
+        )
+        self.client = client or SimpleNamespace(function_invocation_configuration=None)
+        self.messages_received: list[Any] = []
+        self.tools_received: list[Any] | None = None
+        self.last_session: AgentSession | None = None
+
+    @overload
+    def run(
+        self,
+        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        *,
+        stream: Literal[False] = ...,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]]: ...
+
+    @overload
+    def run(
+        self,
+        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        *,
+        stream: Literal[True],
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+    def run(
+        self,
+        messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+        *,
+        stream: bool = False,
+        session: AgentSession | None = None,
+        **kwargs: Any,
+    ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+        if stream:
+
+            async def _stream() -> AsyncIterator[AgentResponseUpdate]:
+                if messages is None:
+                    self.messages_received = []
+                elif isinstance(messages, (str, Content, Message)):
+                    self.messages_received = [messages]
+                else:
+                    self.messages_received = list(messages)
+                self.last_session = session
+                self.tools_received = kwargs.get("tools")
+                for update in self.updates:
+                    yield update
+
+            def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse:
+                return AgentResponse.from_updates(updates)
+
+            return ResponseStream(_stream(), finalizer=_finalize)
+
+        async def _get_response() -> AgentResponse[Any]:
+            return AgentResponse(messages=[], response_id="stub-response")
+
+        return _get_response()
+
+    def create_session(self, **kwargs: Any) -> AgentSession:
+        return AgentSession(session_id=kwargs.get("session_id"))
+
+    def get_session(self, service_session_id: str | ServiceSessionId, *, session_id: str | None = None) -> AgentSession:
+        return AgentSession(session_id=session_id, service_session_id=service_session_id)
+
+
+# Fixtures
+
+
+@pytest.fixture
+def streaming_chat_client_stub() -> type[SupportsChatGetResponse]:
+    """Return the StreamingChatClientStub class for creating test instances."""
+    return StreamingChatClientStub  # type: ignore[return-value]
+
+
+@pytest.fixture
+def stream_from_updates_fixture() -> Callable[[list[ChatResponseUpdate]], StreamFn]:
+    """Return the stream_from_updates helper function."""
+    return stream_from_updates
+
+
+@pytest.fixture
+def stub_agent() -> type[SupportsAgentRun]:
+    """Return the StubAgent class for creating test instances."""
+    return StubAgent  # type: ignore[return-value]
+
+
+# ── Fixtures for golden / integration tests ──
+
+
+@pytest.fixture
+def collect_events() -> Callable[..., Any]:
+    """Return an async helper that collects all events from an async generator."""
+
+    async def _collect(async_gen: AsyncIterable[Any]) -> list[Any]:
+        return [event async for event in async_gen]
+
+    return _collect
+
+
+@pytest.fixture
+def make_agent_wrapper() -> Callable[..., Any]:
+    """Factory that builds an AgentFrameworkAgent from a stream function.
+
+    Usage::
+
+        agent = make_agent_wrapper(
+            stream_fn=stream_from_updates(updates),
+            state_schema=...,
+        )
+        events = [e async for e in agent.run(payload)]
+    """
+    from agent_framework_ag_ui import AgentFrameworkAgent
+
+    def _factory(
+        stream_fn: StreamFn,
+        *,
+        state_schema: Any | None = None,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        require_confirmation: bool = True,
+    ) -> Any:
+        client = StreamingChatClientStub(stream_fn)
+        stub = StubAgent(client=client)
+        return AgentFrameworkAgent(
+            agent=stub,
+            state_schema=state_schema,
+            predict_state_config=predict_state_config,
+            require_confirmation=require_confirmation,
+        )
+
+    return _factory
+
+
+@pytest.fixture
+def make_app() -> Callable[..., Any]:
+    """Factory that builds a FastAPI app with an AG-UI endpoint.
+
+    Usage::
+
+        app = make_app(agent_or_wrapper, path="/test")
+    """
+    from fastapi import FastAPI
+
+    from agent_framework_ag_ui import add_agent_framework_fastapi_endpoint
+
+    def _factory(
+        agent: Any,
+        *,
+        path: str = "/",
+        state_schema: Any | None = None,
+        predict_state_config: dict[str, dict[str, str]] | None = None,
+        default_state: dict[str, Any] | None = None,
+    ) -> FastAPI:
+        app = FastAPI()
+        add_agent_framework_fastapi_endpoint(
+            app,
+            agent,
+            path=path,
+            state_schema=state_schema,
+            predict_state_config=predict_state_config,
+            default_state=default_state,
+        )
+        return app
+
+    return _factory