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

Code before:
    model = GoogleModel('gemini-2.5-flash-image', provider=vertex_provider)

    agent = Agent(model, output_type=BinaryImage)

    result = await agent.run('Generate an image of an axolotl.')
    assert result.output == snapshot(BinaryImage(data=IsBytes(), media_type='image/png', identifier='b037a4'))


async def test_google_httpx_client_is_not_closed(allow_model_requests: None, gemini_api_key: str):
    # This should not raise any errors, see https://github.com/pydantic/pydantic-ai/issues/3242.
    agent = Agent(GoogleModel('gemini-2.5-flash-lite', provider=GoogleProvider(api_key=gemini_api_key)))
    result = await agent.run('What is the capital of France?')
    assert result.output == snapshot('The capital of France is **Paris**.')

    agent = Agent(GoogleModel('gemini-2.5-flash-lite', provider=GoogleProvider(api_key=gemini_api_key)))
    result = await agent.run('What is the capital of Mexico?')
    assert result.output == snapshot('The capital of Mexico is **Mexico City**.')


def test_google_process_response_filters_empty_text_parts(google_provider: GoogleProvider):
    model = GoogleModel('gemini-1.5-flash', provider=google_provider)
    response = _generate_response_with_texts(response_id='resp-123', texts=['', 'first', '', 'second'])

    result = model._process_response(response)

    assert [part.content for part in result.parts if isinstance(part, TextPart)] == ['first', 'second']


async def test_gemini_streamed_response_emits_text_events_for_non_empty_parts():
    chunk = _generate_response_with_texts('stream-1', ['', 'streamed text'])

    async def response_iterator() -> AsyncIterator[GenerateContentResponse]:
        yield chunk

    streamed_response = GeminiStreamedResponse(
        model_request_parameters=ModelRequestParameters(),
        _model_name='gemini-test',
        _response=response_iterator(),
        _timestamp=datetime.datetime.now(datetime.timezone.utc),
        _provider_name='test-provider',
    )

    events = [event async for event in streamed_response._get_event_iterator()]
    assert len(events) == 1
    event = events[0]
    assert isinstance(event, PartStartEvent)
    assert isinstance(event.part, TextPart)
    assert event.part.content == 'streamed text'


def _generate_response_with_texts(response_id: str, texts: list[str]) -> GenerateContentResponse:

Review comment (suggestion):
Small thing, but let's use `assert result.parts == snapshot()`. Then when you run the test, the snapshot will be filled in.

Code after:

    agent = Agent(model, output_type=BinaryImage)

    result = await agent.run('Generate an image of an axolotl.')
    assert result.output == snapshot(BinaryImage(data=IsBytes(), media_type='image/png', identifier='b037a4'))


async def test_google_httpx_client_is_not_closed(allow_model_requests: None, gemini_api_key: str):
    # This should not raise any errors, see https://github.com/pydantic/pydantic-ai/issues/3242.
    agent = Agent(GoogleModel('gemini-2.5-flash-lite', provider=GoogleProvider(api_key=gemini_api_key)))
    result = await agent.run('What is the capital of France?')
    assert result.output == snapshot('The capital of France is **Paris**.')

    agent = Agent(GoogleModel('gemini-2.5-flash-lite', provider=GoogleProvider(api_key=gemini_api_key)))
    result = await agent.run('What is the capital of Mexico?')
    assert result.output == snapshot('The capital of Mexico is **Mexico City**.')


def test_google_process_response_filters_empty_text_parts(google_provider: GoogleProvider):
    model = GoogleModel('gemini-2.5-pro', provider=google_provider)
    response = _generate_response_with_texts(response_id='resp-123', texts=['', 'first', '', 'second'])

    result = model._process_response(response)  # pyright: ignore[reportPrivateUsage]

    assert result.parts == snapshot([TextPart(content='first'), TextPart(content='second')])


async def test_gemini_streamed_response_emits_text_events_for_non_empty_parts():
    chunk = _generate_response_with_texts('stream-1', ['', 'streamed text'])

    async def response_iterator() -> AsyncIterator[GenerateContentResponse]:
        yield chunk

    streamed_response = GeminiStreamedResponse(
        model_request_parameters=ModelRequestParameters(),
        _model_name='gemini-test',
        _response=response_iterator(),
        _timestamp=datetime.datetime.now(datetime.timezone.utc),
        _provider_name='test-provider',
    )

    events = [event async for event in streamed_response._get_event_iterator()]  # pyright: ignore[reportPrivateUsage]
    assert events == snapshot([PartStartEvent(index=0, part=TextPart(content='streamed text'))])


def _generate_response_with_texts(response_id: str, texts: list[str]) -> GenerateContentResponse:
    return GenerateContentResponse.model_validate(
        {
            'response_id': response_id,
            'model_version': 'gemini-test',
            'usage_metadata': GenerateContentResponseUsageMetadata(