1+ """Test for MCP endpoint async generator fix."""
2+
3+ import pytest
4+ from fastapi .testclient import TestClient
5+ from unittest .mock import MagicMock , patch , AsyncMock
6+ from fastapi .responses import StreamingResponse
7+
8+
9+ @pytest .fixture
10+ def client ():
11+ """Create a test client for the FastAPI app."""
12+ from src .ansari .app .main_api import app
13+
14+ return TestClient (app )
15+
16+
17+ @pytest .fixture
18+ def mock_async_presenter ():
19+ """Mock the presenter with async generator to test the fix."""
20+ with patch ("src.ansari.app.main_api.presenter" ) as mock :
21+ # Create a mock async streaming response
22+ async def mock_async_generator ():
23+ """Simulate an async generator like the real presenter."""
24+ yield "This is "
25+ yield "an async "
26+ yield "response"
27+ yield "\n \n **Citations**:\n "
28+ yield "[1] Test Source"
29+
30+ # Create a mock StreamingResponse with async generator
31+ mock_response = MagicMock ()
32+ mock_response .body_iterator = mock_async_generator ()
33+ mock_response .media_type = "text/plain"
34+
35+ def mock_complete (body , message_logger = None ):
36+ return mock_response
37+
38+ mock .complete = mock_complete
39+ yield mock
40+
41+
42+ class TestMCPAsyncFix :
43+ """Test cases for the MCP endpoint async generator fix."""
44+
45+ def test_mcp_endpoint_handles_async_generator (self , client , mock_async_presenter ):
46+ """Test that the MCP endpoint properly handles async generators without errors."""
47+ # Send a request to the MCP endpoint
48+ response = client .post (
49+ "/api/v2/mcp-complete" ,
50+ json = {"messages" : [{"role" : "user" , "content" : "Test async handling" }]},
51+ )
52+
53+ # Should not raise TypeError about async_generator not being iterable
54+ assert response .status_code == 200
55+
56+ # Collect the streamed content
57+ content = response .content
58+
59+ # Verify the original content is present
60+ assert b"This is an async response" in content
61+ assert b"Citations" in content
62+ assert b"[1] Test Source" in content
63+
64+ # Verify the attribution message is added
65+ assert b"ansari.chat" in content
66+ assert b"IT IS ABSOLUTELY CRITICAL" in content
67+
68+ def test_mcp_endpoint_streams_correctly (self , client , mock_async_presenter ):
69+ """Test that the MCP endpoint streams content chunk by chunk."""
70+ response = client .post (
71+ "/api/v2/mcp-complete" ,
72+ json = {"messages" : [{"role" : "user" , "content" : "Test streaming" }]},
73+ )
74+
75+ assert response .status_code == 200
76+
77+ # Get content
78+ full_content = response .content
79+ assert b"This is an async response" in full_content
80+
81+ @patch ("src.ansari.app.main_api.MessageLogger" )
82+ @patch ("src.ansari.app.main_api.db" )
83+ def test_mcp_endpoint_with_real_async_flow (self , mock_db , mock_message_logger , client ):
84+ """Test the complete async flow with proper mocking."""
85+ with patch ("src.ansari.app.main_api.presenter" ) as mock_presenter :
86+ # Create an async generator for testing
87+ async def async_content_generator ():
88+ yield "Hello "
89+ yield "from "
90+ yield "async "
91+ yield "generator"
92+
93+ mock_response = MagicMock ()
94+ mock_response .body_iterator = async_content_generator ()
95+ mock_response .media_type = "text/plain"
96+ mock_presenter .complete .return_value = mock_response
97+
98+ # Make the request
99+ response = client .post (
100+ "/api/v2/mcp-complete" ,
101+ json = {"messages" : [{"role" : "user" , "content" : "Test" }]},
102+ )
103+
104+ assert response .status_code == 200
105+
106+ # Verify content
107+ content = response .content
108+ assert b"Hello from async generator" in content
109+ assert b"ansari.chat" in content
110+
111+ def test_mcp_endpoint_error_handling_with_async (self , client ):
112+ """Test that async errors are handled gracefully."""
113+ with patch ("src.ansari.app.main_api.presenter" ) as mock_presenter :
114+ # Create an async generator that yields successfully
115+ async def working_generator ():
116+ yield "Start "
117+ yield "Middle "
118+ yield "End"
119+
120+ mock_response = MagicMock ()
121+ mock_response .body_iterator = working_generator ()
122+ mock_response .media_type = "text/plain"
123+ mock_presenter .complete .return_value = mock_response
124+
125+ # The request should work correctly
126+ response = client .post (
127+ "/api/v2/mcp-complete" ,
128+ json = {"messages" : [{"role" : "user" , "content" : "Test" }]},
129+ )
130+
131+ # Should succeed with the async generator
132+ assert response .status_code == 200
133+ content = response .content
134+ assert b"Start Middle End" in content
135+ assert b"ansari.chat" in content
0 commit comments