refactor: Remove internal function usage in test_pydantic_ai

Author: harryautomazione

tests/integrations/pydantic_ai/test_pydantic_ai.py

@@ -1637,7 +1637,7 @@ async def test_input_messages_error_handling(sentry_init, capture_events):
     Test that _set_input_messages handles errors gracefully.
     """
     import sentry_sdk
-    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages
+    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span
 
     sentry_init(
         integrations=[PydanticAIIntegration()],
@@ -1646,14 +1646,11 @@ async def test_input_messages_error_handling(sentry_init, capture_events):
     )
 
     with sentry_sdk.start_transaction(op="test", name="test") as transaction:
-        span = sentry_sdk.start_span(op="test_span")
-
         # Pass invalid messages that would cause an error
         invalid_messages = [object()]  # Plain object without expected attributes
 
-        # Should not raise, error is caught internally
-        _set_input_messages(span, invalid_messages)
-
+        # Should not raise, error is caught internally via ai_client_span
+        span = ai_client_span(invalid_messages, None, None, None)
         span.finish()
 
     # Should not crash
@@ -1789,34 +1786,21 @@ async def test_message_parts_with_list_content(sentry_init, capture_events):
     """
     Test that message parts with list content are handled correctly.
     """
-    import sentry_sdk
-    from unittest.mock import MagicMock
-    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages
+    from pydantic_ai import Agent
 
     sentry_init(
         integrations=[PydanticAIIntegration()],
         traces_sample_rate=1.0,
         send_default_pii=True,
     )
 
-    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
-        span = sentry_sdk.start_span(op="test_span")
-
-        # Create message with list content
-        mock_msg = MagicMock()
-        mock_part = MagicMock()
-        mock_part.content = ["item1", "item2", {"complex": "item"}]
-        mock_msg.parts = [mock_part]
-        mock_msg.instructions = None
-
-        messages = [mock_msg]
-
-        # Should handle list content
-        _set_input_messages(span, messages)
+    events = capture_events()
+    agent = Agent("test")
 
-        span.finish()
+    # Run with list content
+    await agent.run(["item1", "item2"])
 
-    # Should not crash
+    (transaction,) = events
     assert transaction is not None
 
 
@@ -1896,35 +1880,27 @@ async def test_message_with_system_prompt_part(sentry_init, capture_events):
     """
     Test that SystemPromptPart is handled with correct role.
     """
-    import sentry_sdk
-    from unittest.mock import MagicMock
-    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages
-    from pydantic_ai import messages
+    from pydantic_ai import Agent, messages
 
     sentry_init(
         integrations=[PydanticAIIntegration()],
         traces_sample_rate=1.0,
         send_default_pii=True,
     )
 
-    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
-        span = sentry_sdk.start_span(op="test_span")
-
-        # Create message with SystemPromptPart
-        system_part = messages.SystemPromptPart(content="You are a helpful assistant")
-
-        mock_msg = MagicMock()
-        mock_msg.parts = [system_part]
-        mock_msg.instructions = None
-
-        msgs = [mock_msg]
+    events = capture_events()
+    agent = Agent("test")
 
-        # Should handle system prompt
-        _set_input_messages(span, msgs)
+    # Create message with SystemPromptPart
+    system_part = messages.SystemPromptPart(content="You are a helpful assistant")
+    history = [
+        messages.ModelRequest(parts=[system_part])
+    ]
 
-        span.finish()
+    # Run with history
+    await agent.run("What did I say?", message_history=history)
 
-    # Should not crash
+    (transaction,) = events
     assert transaction is not None
 
 
@@ -1935,7 +1911,7 @@ async def test_message_with_instructions(sentry_init, capture_events):
     """
     import sentry_sdk
     from unittest.mock import MagicMock
-    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages
+    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import ai_client_span
 
     sentry_init(
         integrations=[PydanticAIIntegration()],
@@ -1944,8 +1920,6 @@ async def test_message_with_instructions(sentry_init, capture_events):
     )
 
     with sentry_sdk.start_transaction(op="test", name="test") as transaction:
-        span = sentry_sdk.start_span(op="test_span")
-
         # Create message with instructions
         mock_msg = MagicMock()
         mock_msg.instructions = "System instructions here"
@@ -1955,9 +1929,8 @@ async def test_message_with_instructions(sentry_init, capture_events):
 
         msgs = [mock_msg]
 
-        # Should extract system prompt from instructions
-        _set_input_messages(span, msgs)
-
+        # Should handle system prompt via ai_client_span wrappers
+        span = ai_client_span(msgs, None, None, None)
         span.finish()
 
     # Should not crash
@@ -1969,25 +1942,21 @@ async def test_set_input_messages_without_prompts(sentry_init, capture_events):
     """
     Test that _set_input_messages respects _should_send_prompts().
     """
-    import sentry_sdk
-    from sentry_sdk.integrations.pydantic_ai.spans.ai_client import _set_input_messages
+    from pydantic_ai import Agent
 
     sentry_init(
         integrations=[PydanticAIIntegration(include_prompts=False)],
         traces_sample_rate=1.0,
         send_default_pii=True,
     )
 
-    with sentry_sdk.start_transaction(op="test", name="test") as transaction:
-        span = sentry_sdk.start_span(op="test_span")
+    events = capture_events()
+    agent = Agent("test")
 
-        # Even with messages, should not set them
-        messages = ["test"]
-        _set_input_messages(span, messages)
+    # Run with prompts disabled
+    await agent.run("test")
 
-        span.finish()
-
-    # Should not crash and should not set messages
+    (transaction,) = events
     assert transaction is not None
 
 
@@ -2705,28 +2674,23 @@ async def test_binary_content_encoding_image(sentry_init, capture_events):
 @pytest.mark.asyncio
 async def test_binary_content_encoding_mixed_content(sentry_init, capture_events):
     """Test that BinaryContent mixed with text content is properly handled."""
+    from pydantic_ai import Agent
+
     sentry_init(
         integrations=[PydanticAIIntegration()],
         traces_sample_rate=1.0,
         send_default_pii=True,
     )
 
     events = capture_events()
+    agent = Agent("test")
 
-    with sentry_sdk.start_transaction(op="test", name="test"):
-        span = sentry_sdk.start_span(op="test_span")
-        binary_content = BinaryContent(
-            data=b"fake_image_bytes", media_type="image/jpeg"
-        )
-        user_part = UserPromptPart(
-            content=["Here is an image:", binary_content, "What do you see?"]
-        )
-        mock_msg = MagicMock()
-        mock_msg.parts = [user_part]
-        mock_msg.instructions = None
+    binary_content = BinaryContent(
+        data=b"fake_image_bytes", media_type="image/jpeg"
+    )
 
-        _set_input_messages(span, [mock_msg])
-        span.finish()
+    # Run with mixed content
+    await agent.run(["Here is an image:", binary_content, "What do you see?"])
 
     (event,) = events
     span_data = event["spans"][0]["data"]