Fixing unit tests

Author: amri369

examples/mcp/caching/main.py

@@ -5,6 +5,8 @@
 import time
 from typing import Any
 
+from mcp.types import ListPromptsResult
+
 from agents import gen_trace_id, trace
 from agents.mcp import MCPServerStreamableHttp
 
@@ -25,8 +27,11 @@ async def run(mcp_server: MCPServerStreamableHttp):
     print("Cached prompts after invoking list_prompts")
     await mcp_server.list_prompts()
     cached_prompts_list = mcp_server._prompts_list
-    for prompt in cached_prompts_list.prompts:
-        print(f"name: {prompt.name}")
+    if isinstance(cached_prompts_list, ListPromptsResult):
+        for prompt in cached_prompts_list.prompts:
+            print(f"name: {prompt.name}")
+    else:
+        print("Failed to cache list_prompts")
 
 async def main():
     async with MCPServerStreamableHttp(

tests/mcp/test_caching.py

@@ -84,32 +84,33 @@ async def test_server_caching_prompts_works(
         Prompt(name="prompt2"),
     ]
 
-    mock_list_prompts.return_value = ListPromptsResult(prompts=prompts)
+    list_prompts = ListPromptsResult(prompts=prompts)
+    mock_list_prompts.return_value = list_prompts
 
     async with server:
 
         # Call list_prompts() multiple times
         result_prompts = await server.list_prompts()
-        assert result_prompts == prompts
+        assert result_prompts == list_prompts
 
         assert mock_list_prompts.call_count == 1, "list_prompts() should have been called once"
 
         # Call list_prompts() again, should return the cached value
         result_prompts = await server.list_prompts()
-        assert result_prompts == prompts
+        assert result_prompts == list_prompts
 
         assert mock_list_prompts.call_count == 1, ("list_prompts() "
                                                    "should not have been called again")
 
         # Invalidate the cache and call list_prompts() again
         server.invalidate_prompts_cache()
         result_prompts = await server.list_prompts()
-        assert result_prompts == prompts
+        assert result_prompts == list_prompts
 
         assert mock_list_prompts.call_count == 2, ("list_prompts() "
                                                    "should be called again")
 
         # Without invalidating the cache, calling list_prompts()
         # again should return the cached value
         result_prompts = await server.list_prompts()
-        assert result_prompts == prompts
+        assert result_prompts == list_prompts