From 58985f9d5712d55832021aac359eacf953ddf8dd Mon Sep 17 00:00:00 2001
From: Mery-Sanz <maria@aliasrobotics.com>
Date: Fri, 11 Apr 2025 08:24:26 +0200
Subject: [PATCH] fix mcp test

---
 tests/helpers.py                     |  2 +-
 tests/mcp/helpers.py                 | 21 +++++++++++++++++++++
 tests/mcp/test_caching.py            |  2 +-
 tests/mcp/test_connect_disconnect.py |  2 +-
 tests/mcp/test_mcp_tracing.py        |  2 +-
 tests/mcp/test_mcp_util.py           |  2 +-
 tests/mcp/test_runner_calls_mcp.py   |  2 +-
 tests/voice/helpers.py               |  2 +-
 8 files changed, 28 insertions(+), 7 deletions(-)
 create mode 100644 tests/mcp/helpers.py

diff --git a/tests/helpers.py b/tests/helpers.py
index 8ff153c1..26db6791 100644
--- a/tests/helpers.py
+++ b/tests/helpers.py
@@ -5,7 +5,7 @@ from typing import Any
 from mcp import Tool as MCPTool
 from mcp.types import CallToolResult, TextContent
 
-from agents.mcp import MCPServer
+from cai.sdk.agents.mcp import MCPServer
 
 tee = shutil.which("tee") or ""
 assert tee, "tee not found"
diff --git a/tests/mcp/helpers.py b/tests/mcp/helpers.py
new file mode 100644
index 00000000..98b0202c
--- /dev/null
+++ b/tests/mcp/helpers.py
@@ -0,0 +1,21 @@
+try:
+    from cai.sdk.agents.voice import StreamedAudioResult
+except ImportError:
+    pass
+
+
+async def extract_events(result: StreamedAudioResult) -> tuple[list[str], list[bytes]]:
+    """Collapse pipeline stream events to simple labels for ordering assertions."""
+    flattened: list[str] = []
+    audio_chunks: list[bytes] = []
+
+    async for ev in result.stream():
+        if ev.type == "voice_stream_event_audio":
+            if ev.data is not None:
+                audio_chunks.append(ev.data.tobytes())
+            flattened.append("audio")
+        elif ev.type == "voice_stream_event_lifecycle":
+            flattened.append(ev.event)
+        elif ev.type == "voice_stream_event_error":
+            flattened.append("error")
+    return flattened, audio_chunks
diff --git a/tests/mcp/test_caching.py b/tests/mcp/test_caching.py
index 16dcd227..ee53640e 100644
--- a/tests/mcp/test_caching.py
+++ b/tests/mcp/test_caching.py
@@ -5,7 +5,7 @@ from mcp.types import ListToolsResult, Tool as MCPTool
 
 from cai.sdk.agents.mcp import MCPServerStdio
 
-from tests.mcp.helpers import DummyStreamsContextManager, tee
+from tests.helpers import DummyStreamsContextManager, tee
 
 
 @pytest.mark.asyncio
diff --git a/tests/mcp/test_connect_disconnect.py b/tests/mcp/test_connect_disconnect.py
index 85ee6ec6..7565eddb 100644
--- a/tests/mcp/test_connect_disconnect.py
+++ b/tests/mcp/test_connect_disconnect.py
@@ -5,7 +5,7 @@ from mcp.types import ListToolsResult, Tool as MCPTool
 
 from cai.sdk.agents.mcp import MCPServerStdio
 
-from .helpers import DummyStreamsContextManager, tee
+from tests.helpers import DummyStreamsContextManager, tee
 
 
 @pytest.mark.asyncio
diff --git a/tests/mcp/test_mcp_tracing.py b/tests/mcp/test_mcp_tracing.py
index a4ebc653..2c832e56 100644
--- a/tests/mcp/test_mcp_tracing.py
+++ b/tests/mcp/test_mcp_tracing.py
@@ -6,7 +6,7 @@ from cai.sdk.agents import Agent, Runner
 from tests.fake_model import FakeModel
 from tests.core.test_responses import get_function_tool, get_function_tool_call, get_text_message
 from tests.testing_processor import SPAN_PROCESSOR_TESTING, fetch_normalized_spans
-from .helpers import FakeMCPServer
+from tests.helpers import FakeMCPServer
 
 
 @pytest.mark.asyncio
diff --git a/tests/mcp/test_mcp_util.py b/tests/mcp/test_mcp_util.py
index f62d3ab0..5e8d6da8 100644
--- a/tests/mcp/test_mcp_util.py
+++ b/tests/mcp/test_mcp_util.py
@@ -9,7 +9,7 @@ from cai.sdk.agents import FunctionTool, RunContextWrapper
 from cai.sdk.agents.exceptions import AgentsException, ModelBehaviorError
 from cai.sdk.agents.mcp import MCPServer, MCPUtil
 
-from .helpers import FakeMCPServer
+from tests.helpers import FakeMCPServer
 
 
 class Foo(BaseModel):
diff --git a/tests/mcp/test_runner_calls_mcp.py b/tests/mcp/test_runner_calls_mcp.py
index c879e8fd..35eefd6b 100644
--- a/tests/mcp/test_runner_calls_mcp.py
+++ b/tests/mcp/test_runner_calls_mcp.py
@@ -7,7 +7,7 @@ from cai.sdk.agents import Agent, ModelBehaviorError, Runner, UserError
 
 from tests.fake_model import FakeModel
 from tests.core.test_responses import get_function_tool_call, get_text_message
-from .helpers import FakeMCPServer
+from tests.helpers import FakeMCPServer
 
 
 @pytest.mark.asyncio
diff --git a/tests/voice/helpers.py b/tests/voice/helpers.py
index ae902dc1..98b0202c 100644
--- a/tests/voice/helpers.py
+++ b/tests/voice/helpers.py
@@ -1,5 +1,5 @@
 try:
-    from agents.voice import StreamedAudioResult
+    from cai.sdk.agents.voice import StreamedAudioResult
 except ImportError:
     pass