Repository: pydantic/pydantic-ai
File: tests/test_ag_ui.py (Python)

Code before:
    assert isinstance(result, AudioUrl)
    assert result.url == 'http://example.com/audio.mp3'
    assert result.media_type == 'audio/mpeg'
    assert result.identifier == 'aud1'

    # Test DocumentUrl
    doc_url_part = BinaryInputContent(url='http://example.com/doc.pdf', mime_type='application/pdf', id='doc1')
    result = AGUIAdapter.load_binary_part(doc_url_part)
    assert isinstance(result, DocumentUrl)
    assert result.url == 'http://example.com/doc.pdf'
    assert result.media_type == 'application/pdf'
    assert result.identifier == 'doc1'

    # Test data field
    data_part = BinaryInputContent(data='SGVsbG8gd29ybGQ=', mime_type='text/plain')
    result = AGUIAdapter.load_binary_part(data_part)
    assert isinstance(result, BinaryContent)
    assert result.data == b'Hello world'
    assert result.media_type == 'text/plain'

    # Test ValueError
    with pytest.raises(ValueError, match='BinaryInputContent must have either a `url` or `data` field.'):
        AGUIAdapter.load_binary_part(BinaryInputContent(id='some_id', mime_type='text/plain'))


def test_add_assistant_tool_parts() -> None:
    """Test the _add_assistant_tool_parts method of the AGUIAdapter."""
    # Case 1: Regular tool call
    builder = MessagesBuilder()
    tool_calls_map: dict[str, str] = {}
    regular_tool_call = ToolCall(
        id='regular_call_1',
        type='function',
        function=FunctionCall(name='my_tool', arguments='{"arg": "value"}'),
    )
    tool_calls_list = [regular_tool_call]

    AGUIAdapter.add_assistant_tool_parts(builder, tool_calls_list, tool_calls_map)

    assert len(builder.messages) == 1
    assert len(builder.messages[0].parts) == 1
    part = builder.messages[0].parts[0]
    assert isinstance(part, ToolCallPart)
    assert part.tool_name == 'my_tool'
    assert part.tool_call_id == 'regular_call_1'
    assert part.args == '{"arg": "value"}'
    assert tool_calls_map == {'regular_call_1': 'my_tool'}

    # Case 2: Built-in tool call
    builder = MessagesBuilder()
    tool_calls_map = {}

Review comment (refactor):
We shouldn't need to test these helper methods directly, as all of these features have already been tested through the existing tests, except for the new binary content. So I'd expect for the _only_ new tests to be related to those binary parts.

If these tests were necessary because of the refactoring of the original massive method into separate ones, I'd prefer to revert that and have a massive method again, which was 100% test-covered.

Code after:
                    TextPart(content='Assistant message'),
                    ToolCallPart(tool_name='tool_call_1', args='{}', tool_call_id='tool_call_1'),
                    ToolCallPart(tool_name='tool_call_2', args='{}', tool_call_id='tool_call_2'),
                ],
                timestamp=IsDatetime(),
            ),
            ModelRequest(
                parts=[
                    ToolReturnPart(
                        tool_name='tool_call_1',
                        content='Tool message',
                        tool_call_id='tool_call_1',
                        timestamp=IsDatetime(),
                    ),
                    ToolReturnPart(
                        tool_name='tool_call_2',
                        content='Tool message',
                        tool_call_id='tool_call_2',
                        timestamp=IsDatetime(),
                    ),
                    UserPromptPart(
                        content='User message',
                        timestamp=IsDatetime(),
                    ),
                ]
            ),
            ModelResponse(
                parts=[TextPart(content='Assistant message')],
                timestamp=IsDatetime(),
            ),
        ]
    )


async def test_builtin_tool_call() -> None:
    async def stream_function(
        messages: list[ModelMessage], agent_info: AgentInfo
    ) -> AsyncIterator[BuiltinToolCallsReturns | DeltaToolCalls | str]:
        yield {
            0: BuiltinToolCallPart(
                tool_name=WebSearchTool.kind,
                args='{"query":',
                tool_call_id='search_1',
                provider_name='function',
            )
        }
        yield {
            0: DeltaToolCall(
                json_args='"Hello world"}',
                tool_call_id='search_1',
            )