From 50ee4a400589a60086f865d8dfb68ea9b3a35125 Mon Sep 17 00:00:00 2001
From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
Date: Tue, 16 Sep 2025 23:39:36 -0700
Subject: [PATCH] Added unit tests

---
 .../main/tests/main/test_middleware.py        | 582 +++++++++++++++++-
 1 file changed, 579 insertions(+), 3 deletions(-)

diff --git a/python/packages/main/tests/main/test_middleware.py b/python/packages/main/tests/main/test_middleware.py
index 454332c5f8..b49b473092 100644
--- a/python/packages/main/tests/main/test_middleware.py
+++ b/python/packages/main/tests/main/test_middleware.py
@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from collections.abc import AsyncIterable, Awaitable, Callable
+from collections.abc import AsyncIterable, Awaitable, Callable, MutableSequence
 from typing import Any
 from unittest.mock import MagicMock
 
@@ -11,10 +11,16 @@ from agent_framework import (
     AgentProtocol,
     AgentRunResponse,
     AgentRunResponseUpdate,
+    ChatAgent,
     ChatMessage,
+    ChatResponse,
+    ChatResponseUpdate,
+    FunctionCallContent,
+    FunctionResultContent,
     Role,
     TextContent,
 )
+from agent_framework._clients import BaseChatClient
 from agent_framework._middleware import (
     AgentMiddleware,
     AgentMiddlewarePipeline,
@@ -23,7 +29,8 @@ from agent_framework._middleware import (
     FunctionMiddleware,
     FunctionMiddlewarePipeline,
 )
-from agent_framework._tools import AIFunction
+from agent_framework._tools import AIFunction, use_function_invocation
+from agent_framework._types import ChatOptions
 
 
 class TestAgentRunContext:
@@ -715,7 +722,7 @@ class TestStreamingScenarios:
         ]
 
 
-# Helper classes and fixtures
+# region Helper classes and fixtures
 
 
 class FunctionTestArgs(BaseModel):
@@ -760,3 +767,572 @@ def mock_function() -> AIFunction[Any, Any]:
     function = MagicMock(spec=AIFunction[Any, Any])
     function.name = "test_function"
     return function
+
+
+@use_function_invocation
+class MockChatClient(BaseChatClient):
+    """Mock chat client for ChatAgent integration tests."""
+
+    call_count: int = Field(default=0)
+    responses: list[ChatResponse] = Field(default_factory=lambda: [])
+    streaming_responses: list[list[ChatResponseUpdate]] = Field(default_factory=lambda: [])
+
+    async def _inner_get_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> ChatResponse:
+        """Return a mock response."""
+        self.call_count += 1
+        if self.responses:
+            return self.responses.pop(0)
+        return ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Mock response")])
+
+    async def _inner_get_streaming_response(
+        self,
+        *,
+        messages: MutableSequence[ChatMessage],
+        chat_options: ChatOptions,
+        **kwargs: Any,
+    ) -> AsyncIterable[ChatResponseUpdate]:
+        """Return mock streaming responses."""
+        self.call_count += 1
+        if self.streaming_responses:
+            for update in self.streaming_responses.pop(0):
+                yield update
+        else:
+            yield ChatResponseUpdate(contents=[TextContent(text="Mock")], role=Role.ASSISTANT)
+            yield ChatResponseUpdate(contents=[TextContent(text=" streaming response")], role=Role.ASSISTANT)
+
+    def service_url(self) -> str:
+        return "https://mock.example.com"
+
+
+@pytest.fixture
+def mock_chat_client() -> MockChatClient:
+    """Mock chat client fixture."""
+    return MockChatClient()
+
+
+# region ChatAgent Tests
+
+
+class TestChatAgentClassBasedMiddleware:
+    """Test cases for class-based middleware integration with ChatAgent."""
+
+    async def test_class_based_agent_middleware_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
+        """Test class-based agent middleware with ChatAgent."""
+        execution_order: list[str] = []
+
+        class TrackingAgentMiddleware(AgentMiddleware):
+            def __init__(self, name: str):
+                self.name = name
+
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                execution_order.append(f"{self.name}_before")
+                await next(context)
+                execution_order.append(f"{self.name}_after")
+
+        # Create ChatAgent with middleware
+        middleware = TrackingAgentMiddleware("agent_middleware")
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert response.messages[0].role == Role.ASSISTANT
+        assert response.messages[0].text == "Mock response"
+        assert mock_chat_client.call_count == 1
+
+        # Verify middleware execution order
+        assert execution_order == ["agent_middleware_before", "agent_middleware_after"]
+
+    async def test_class_based_function_middleware_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
+        """Test class-based function middleware with ChatAgent."""
+        execution_order: list[str] = []
+
+        class TrackingFunctionMiddleware(FunctionMiddleware):
+            def __init__(self, name: str):
+                self.name = name
+
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                next: Callable[[FunctionInvocationContext], Awaitable[None]],
+            ) -> None:
+                execution_order.append(f"{self.name}_before")
+                await next(context)
+                execution_order.append(f"{self.name}_after")
+
+        # Create ChatAgent with function middleware (no tools, so function middleware won't be triggered)
+        middleware = TrackingFunctionMiddleware("function_middleware")
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert mock_chat_client.call_count == 1
+
+        # Note: Function middleware won't execute since no function calls are made
+        assert execution_order == []
+
+
+class TestChatAgentFunctionBasedMiddleware:
+    """Test cases for function-based middleware integration with ChatAgent."""
+
+    async def test_function_based_agent_middleware_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
+        """Test function-based agent middleware with ChatAgent."""
+        execution_order: list[str] = []
+
+        async def tracking_agent_middleware(
+            context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+        ) -> None:
+            execution_order.append("agent_function_before")
+            await next(context)
+            execution_order.append("agent_function_after")
+
+        # Create ChatAgent with function middleware
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[tracking_agent_middleware])
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert response.messages[0].role == Role.ASSISTANT
+        assert response.messages[0].text == "Mock response"
+        assert mock_chat_client.call_count == 1
+
+        # Verify middleware execution order
+        assert execution_order == ["agent_function_before", "agent_function_after"]
+
+    async def test_function_based_function_middleware_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
+        """Test function-based function middleware with ChatAgent."""
+        execution_order: list[str] = []
+
+        async def tracking_function_middleware(
+            context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
+        ) -> None:
+            execution_order.append("function_function_before")
+            await next(context)
+            execution_order.append("function_function_after")
+
+        # Create ChatAgent with function middleware (no tools, so function middleware won't be triggered)
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[tracking_function_middleware])
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert mock_chat_client.call_count == 1
+
+        # Note: Function middleware won't execute since no function calls are made
+        assert execution_order == []
+
+
+class TestChatAgentStreamingMiddleware:
+    """Test cases for streaming middleware integration with ChatAgent."""
+
+    async def test_agent_middleware_with_streaming(self, mock_chat_client: MockChatClient) -> None:
+        """Test agent middleware with streaming ChatAgent responses."""
+        execution_order: list[str] = []
+        streaming_flags: list[bool] = []
+
+        class StreamingTrackingMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                execution_order.append("middleware_before")
+                streaming_flags.append(context.is_streaming)
+                await next(context)
+                execution_order.append("middleware_after")
+
+        # Create ChatAgent with middleware
+        middleware = StreamingTrackingMiddleware()
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
+
+        # Set up mock streaming responses
+        mock_chat_client.streaming_responses = [
+            [
+                ChatResponseUpdate(contents=[TextContent(text="Streaming")], role=Role.ASSISTANT),
+                ChatResponseUpdate(contents=[TextContent(text=" response")], role=Role.ASSISTANT),
+            ]
+        ]
+
+        # Execute streaming
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        updates: list[AgentRunResponseUpdate] = []
+        async for update in agent.run_stream(messages):
+            updates.append(update)
+
+        # Verify streaming response
+        assert len(updates) == 2
+        assert updates[0].text == "Streaming"
+        assert updates[1].text == " response"
+        assert mock_chat_client.call_count == 1
+
+        # Verify middleware was called and streaming flag was set correctly
+        assert execution_order == ["middleware_before", "middleware_after"]
+        assert streaming_flags == [True]  # Context should indicate streaming
+
+    async def test_non_streaming_vs_streaming_flag_validation(self, mock_chat_client: MockChatClient) -> None:
+        """Test that is_streaming flag is correctly set for different execution modes."""
+        streaming_flags: list[bool] = []
+
+        class FlagTrackingMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                streaming_flags.append(context.is_streaming)
+                await next(context)
+
+        # Create ChatAgent with middleware
+        middleware = FlagTrackingMiddleware()
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+
+        # Test non-streaming execution
+        response = await agent.run(messages)
+        assert response is not None
+
+        # Test streaming execution
+        async for _ in agent.run_stream(messages):
+            pass
+
+        # Verify flags: [non-streaming, streaming]
+        assert streaming_flags == [False, True]
+
+
+class TestChatAgentMultipleMiddlewareOrdering:
+    """Test cases for multiple middleware execution order with ChatAgent."""
+
+    async def test_multiple_agent_middleware_execution_order(self, mock_chat_client: MockChatClient) -> None:
+        """Test that multiple agent middlewares execute in correct order with ChatAgent."""
+        execution_order: list[str] = []
+
+        class OrderedMiddleware(AgentMiddleware):
+            def __init__(self, name: str):
+                self.name = name
+
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                execution_order.append(f"{self.name}_before")
+                await next(context)
+                execution_order.append(f"{self.name}_after")
+
+        # Create multiple middlewares
+        middleware1 = OrderedMiddleware("first")
+        middleware2 = OrderedMiddleware("second")
+        middleware3 = OrderedMiddleware("third")
+
+        # Create ChatAgent with multiple middlewares
+        agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware1, middleware2, middleware3])
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert mock_chat_client.call_count == 1
+
+        # Verify execution order (should be nested: first wraps second wraps third)
+        expected_order = ["first_before", "second_before", "third_before", "third_after", "second_after", "first_after"]
+        assert execution_order == expected_order
+
+    async def test_mixed_middleware_types_with_chat_agent(self, mock_chat_client: MockChatClient) -> None:
+        """Test mixed class and function-based middlewares with ChatAgent."""
+        execution_order: list[str] = []
+
+        class ClassAgentMiddleware(AgentMiddleware):
+            async def process(
+                self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+            ) -> None:
+                execution_order.append("class_agent_before")
+                await next(context)
+                execution_order.append("class_agent_after")
+
+        async def function_agent_middleware(
+            context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
+        ) -> None:
+            execution_order.append("function_agent_before")
+            await next(context)
+            execution_order.append("function_agent_after")
+
+        class ClassFunctionMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                next: Callable[[FunctionInvocationContext], Awaitable[None]],
+            ) -> None:
+                execution_order.append("class_function_before")
+                await next(context)
+                execution_order.append("class_function_after")
+
+        async def function_function_middleware(
+            context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
+        ) -> None:
+            execution_order.append("function_function_before")
+            await next(context)
+            execution_order.append("function_function_after")
+
+        # Create ChatAgent with mixed middleware types (no tools, focusing on agent middleware)
+        agent = ChatAgent(
+            chat_client=mock_chat_client,
+            middleware=[
+                ClassAgentMiddleware(),
+                function_agent_middleware,
+                ClassFunctionMiddleware(),  # Won't execute without function calls
+                function_function_middleware,  # Won't execute without function calls
+            ],
+        )
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="test message")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert mock_chat_client.call_count == 1
+
+        # Verify that agent middlewares were executed in correct order
+        # (Function middlewares won't execute since no functions are called)
+        expected_order = ["class_agent_before", "function_agent_before", "function_agent_after", "class_agent_after"]
+        assert execution_order == expected_order
+
+
+# region Tool Functions for Testing
+
+
+def sample_tool_function(location: str) -> str:
+    """A simple tool function for middleware testing."""
+    return f"Weather in {location}: sunny"
+
+
+# region ChatAgent Function Middleware Tests with Tools
+
+
+class TestChatAgentFunctionMiddlewareWithTools:
+    """Test cases for function middleware integration with ChatAgent when tools are used."""
+
+    async def test_class_based_function_middleware_with_tool_calls(self, mock_chat_client: MockChatClient) -> None:
+        """Test class-based function middleware with ChatAgent when function calls are made."""
+        execution_order: list[str] = []
+
+        class TrackingFunctionMiddleware(FunctionMiddleware):
+            def __init__(self, name: str):
+                self.name = name
+
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                next: Callable[[FunctionInvocationContext], Awaitable[None]],
+            ) -> None:
+                execution_order.append(f"{self.name}_before")
+                await next(context)
+                execution_order.append(f"{self.name}_after")
+
+        # Set up mock to return a function call first, then a regular response
+        function_call_response = ChatResponse(
+            messages=[
+                ChatMessage(
+                    role=Role.ASSISTANT,
+                    contents=[
+                        FunctionCallContent(
+                            call_id="call_123",
+                            name="sample_tool_function",
+                            arguments='{"location": "Seattle"}',
+                        )
+                    ],
+                )
+            ]
+        )
+        final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
+
+        mock_chat_client.responses = [function_call_response, final_response]
+
+        # Create ChatAgent with function middleware and tools
+        middleware = TrackingFunctionMiddleware("function_middleware")
+        agent = ChatAgent(
+            chat_client=mock_chat_client,
+            middleware=[middleware],
+            tools=[sample_tool_function],
+        )
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="Get weather for Seattle")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert mock_chat_client.call_count == 2  # Two calls: one for function call, one for final response
+
+        # Verify function middleware was executed
+        assert execution_order == ["function_middleware_before", "function_middleware_after"]
+
+        # Verify function call and result are in the response
+        all_contents = [content for message in response.messages for content in message.contents]
+        function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
+        function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
+
+        assert len(function_calls) == 1
+        assert len(function_results) == 1
+        assert function_calls[0].name == "sample_tool_function"
+        assert function_results[0].call_id == function_calls[0].call_id
+
+    async def test_function_based_function_middleware_with_tool_calls(self, mock_chat_client: MockChatClient) -> None:
+        """Test function-based function middleware with ChatAgent when function calls are made."""
+        execution_order: list[str] = []
+
+        async def tracking_function_middleware(
+            context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
+        ) -> None:
+            execution_order.append("function_middleware_before")
+            await next(context)
+            execution_order.append("function_middleware_after")
+
+        # Set up mock to return a function call first, then a regular response
+        function_call_response = ChatResponse(
+            messages=[
+                ChatMessage(
+                    role=Role.ASSISTANT,
+                    contents=[
+                        FunctionCallContent(
+                            call_id="call_456",
+                            name="sample_tool_function",
+                            arguments='{"location": "San Francisco"}',
+                        )
+                    ],
+                )
+            ]
+        )
+        final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
+
+        mock_chat_client.responses = [function_call_response, final_response]
+
+        # Create ChatAgent with function middleware and tools
+        agent = ChatAgent(
+            chat_client=mock_chat_client,
+            middleware=[tracking_function_middleware],
+            tools=[sample_tool_function],
+        )
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="Get weather for San Francisco")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert mock_chat_client.call_count == 2  # Two calls: one for function call, one for final response
+
+        # Verify function middleware was executed
+        assert execution_order == ["function_middleware_before", "function_middleware_after"]
+
+        # Verify function call and result are in the response
+        all_contents = [content for message in response.messages for content in message.contents]
+        function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
+        function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
+
+        assert len(function_calls) == 1
+        assert len(function_results) == 1
+        assert function_calls[0].name == "sample_tool_function"
+        assert function_results[0].call_id == function_calls[0].call_id
+
+    async def test_mixed_agent_and_function_middleware_with_tool_calls(self, mock_chat_client: MockChatClient) -> None:
+        """Test both agent and function middleware with ChatAgent when function calls are made."""
+        execution_order: list[str] = []
+
+        class TrackingAgentMiddleware(AgentMiddleware):
+            async def process(
+                self,
+                context: AgentRunContext,
+                next: Callable[[AgentRunContext], Awaitable[None]],
+            ) -> None:
+                execution_order.append("agent_middleware_before")
+                await next(context)
+                execution_order.append("agent_middleware_after")
+
+        class TrackingFunctionMiddleware(FunctionMiddleware):
+            async def process(
+                self,
+                context: FunctionInvocationContext,
+                next: Callable[[FunctionInvocationContext], Awaitable[None]],
+            ) -> None:
+                execution_order.append("function_middleware_before")
+                await next(context)
+                execution_order.append("function_middleware_after")
+
+        # Set up mock to return a function call first, then a regular response
+        function_call_response = ChatResponse(
+            messages=[
+                ChatMessage(
+                    role=Role.ASSISTANT,
+                    contents=[
+                        FunctionCallContent(
+                            call_id="call_789",
+                            name="sample_tool_function",
+                            arguments='{"location": "New York"}',
+                        )
+                    ],
+                )
+            ]
+        )
+        final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
+
+        mock_chat_client.responses = [function_call_response, final_response]
+
+        # Create ChatAgent with both agent and function middleware and tools
+        agent = ChatAgent(
+            chat_client=mock_chat_client,
+            middleware=[TrackingAgentMiddleware(), TrackingFunctionMiddleware()],
+            tools=[sample_tool_function],
+        )
+
+        # Execute the agent
+        messages = [ChatMessage(role=Role.USER, text="Get weather for New York")]
+        response = await agent.run(messages)
+
+        # Verify response
+        assert response is not None
+        assert len(response.messages) > 0
+        assert mock_chat_client.call_count == 2  # Two calls: one for function call, one for final response
+
+        # Verify middleware execution order: agent middleware wraps everything,
+        # function middleware only for function calls
+        expected_order = [
+            "agent_middleware_before",
+            "function_middleware_before",
+            "function_middleware_after",
+            "agent_middleware_after",
+        ]
+        assert execution_order == expected_order
+
+        # Verify function call and result are in the response
+        all_contents = [content for message in response.messages for content in message.contents]
+        function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
+        function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
+
+        assert len(function_calls) == 1
+        assert len(function_results) == 1
+        assert function_calls[0].name == "sample_tool_function"
+        assert function_results[0].call_id == function_calls[0].call_id