patch further down call stack

Author: alexander-alderman-webb

tests/integrations/openai_agents/test_openai_agents.py

@@ -1230,7 +1230,6 @@ async def test_hosted_mcp_tool_propagation_header_streamed(
     )
 
     client = AsyncOpenAI(api_key="z")
-    client.responses._post = AsyncMock(return_value=EXAMPLE_RESPONSE)
 
     model = OpenAIResponsesModel(model="gpt-4", openai_client=client)
 
@@ -1245,8 +1244,14 @@ async def test_hosted_mcp_tool_propagation_header_streamed(
         release="d08ebdb9309e1b004c6f52202de58a09c2268e42",
     )
 
+    request = httpx.Request(
+        "POST",
+        "/responses",
+    )
+
     response = httpx.Response(
         200,
+        request=request,
         content=async_iterator(
             sse_chunks(
                 [
@@ -1308,22 +1313,10 @@ async def test_hosted_mcp_tool_propagation_header_streamed(
         ),
     )
 
-    # Emulate https://github.com/openai/openai-python/blob/656e3cab4a18262a49b961d41293367e45ee71b9/src/openai/_base_client.py#L1751
-    api_response = AsyncAPIResponse(
-        raw=response,
-        cast_to=Response,
-        client=client,
-        stream=True,
-        stream_cls=AsyncStream[ResponseStreamEvent],
-        options=FinalRequestOptions.construct(method="post", url="/responses"),
-        retries_taken=0,
-    )
-
     with patch.object(
-        model._client.responses,
-        "create",
-        # based on https://github.com/openai/openai-python/blob/656e3cab4a18262a49b961d41293367e45ee71b9/src/openai/_base_client.py#L1763
-        return_value=await api_response.parse(),
+        agent_with_tool.model._client._client,
+        "send",
+        return_value=response,
     ) as create, mock.patch(
         "sentry_sdk.tracing_utils.Random.randrange", return_value=500000
     ):
@@ -1341,13 +1334,17 @@ async def test_hosted_mcp_tool_propagation_header_streamed(
             async for event in result.stream_events():
                 pass
 
-            ai_client_span = transaction._span_recorder.spans[-1]
+            ai_client_span = next(
+                span
+                for span in transaction._span_recorder.spans
+                if span.op == "gen_ai.chat"
+            )
 
         args, kwargs = create.call_args
 
-        assert "tools" in kwargs
-        assert len(kwargs["tools"]) == 1
-        hosted_mcp_tool = kwargs["tools"][0]
+        request = args[0]
+        body = json.loads(request.content.decode("utf-8"))
+        hosted_mcp_tool = body["tools"][0]
 
         assert hosted_mcp_tool["headers"][
             "sentry-trace"
@@ -2818,7 +2815,6 @@ async def test_streaming_ttft_on_chat_span(sentry_init, test_agent, async_iterat
     should NOT trigger TTFT.
     """
     client = AsyncOpenAI(api_key="z")
-    client.responses._post = AsyncMock(return_value=EXAMPLE_RESPONSE)
 
     model = OpenAIResponsesModel(model="gpt-4", openai_client=client)
 
@@ -2831,8 +2827,14 @@ async def test_streaming_ttft_on_chat_span(sentry_init, test_agent, async_iterat
         traces_sample_rate=1.0,
     )
 
+    request = httpx.Request(
+        "POST",
+        "/responses",
+    )
+
     response = httpx.Response(
         200,
+        request=request,
         content=async_iterator(
             sse_chunks(
                 [
@@ -2912,22 +2914,11 @@ async def test_streaming_ttft_on_chat_span(sentry_init, test_agent, async_iterat
         ),
     )
 
-    # Emulate https://github.com/openai/openai-python/blob/656e3cab4a18262a49b961d41293367e45ee71b9/src/openai/_base_client.py#L1751
-    api_response = AsyncAPIResponse(
-        raw=response,
-        cast_to=Response,
-        client=client,
-        stream=True,
-        stream_cls=AsyncStream[ResponseStreamEvent],
-        options=FinalRequestOptions.construct(method="post", url="/responses"),
-        retries_taken=0,
-    )
-
+    # Patching https://github.com/openai/openai-python/blob/656e3cab4a18262a49b961d41293367e45ee71b9/src/openai/_base_client.py#L1604
     with patch.object(
-        model._client.responses,
-        "create",
-        # based on https://github.com/openai/openai-python/blob/656e3cab4a18262a49b961d41293367e45ee71b9/src/openai/_base_client.py#L1763
-        return_value=await api_response.parse(),
+        agent_with_tool.model._client._client,
+        "send",
+        return_value=response,
     ) as _:
         with sentry_sdk.start_transaction(
             name="test_ttft", sampled=True