diff --git a/tests/test_agent_runner.py b/tests/test_agent_runner.py
index ce0c5804..b09fdc68 100644
--- a/tests/test_agent_runner.py
+++ b/tests/test_agent_runner.py
@@ -6,7 +6,7 @@ from typing import Any
 import pytest
 from typing_extensions import TypedDict
 
-from agents import (
+from cai.sdk.agents import (
     Agent,
     GuardrailFunctionOutput,
     Handoff,
@@ -23,8 +23,8 @@ from agents import (
     UserError,
     handoff,
 )
-from agents.agent import ToolsToFinalOutputResult
-from agents.tool import FunctionToolResult, function_tool
+from cai.sdk.agents.agent import ToolsToFinalOutputResult
+from cai.sdk.agents.tool import FunctionToolResult, function_tool
 
 from .fake_model import FakeModel
 from .test_responses import (
diff --git a/tests/test_agent_runner_streamed.py b/tests/test_agent_runner_streamed.py
index 87a76a70..8d84126b 100644
--- a/tests/test_agent_runner_streamed.py
+++ b/tests/test_agent_runner_streamed.py
@@ -6,7 +6,7 @@ from typing import Any
 import pytest
 from typing_extensions import TypedDict
 
-from agents import (
+from cai.sdk.agents import (
     Agent,
     GuardrailFunctionOutput,
     Handoff,
@@ -20,9 +20,9 @@ from agents import (
     UserError,
     handoff,
 )
-from agents.items import RunItem
-from agents.run import RunConfig
-from agents.stream_events import AgentUpdatedStreamEvent
+from cai.sdk.agents.items import RunItem
+from cai.sdk.agents.run import RunConfig
+from cai.sdk.agents.stream_events import AgentUpdatedStreamEvent
 
 from .fake_model import FakeModel
 from .test_responses import (