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

Code before:

    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:
    return GenerateContentResponse.model_validate(
        {
            'response_id': response_id,
            'model_version': 'gemini-test',
            'usage_metadata': GenerateContentResponseUsageMetadata(
                prompt_token_count=0,
                candidates_token_count=0,
            ),
            'candidates': [
                {
                    'finish_reason': GoogleFinishReason.STOP,
                    'content': {
                        'role': 'model',
                        'parts': [{'text': text} for text in texts],
                    },
                }
            ],
        }
    )


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

Code after:
    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(
                prompt_token_count=0,
                candidates_token_count=0,
            ),
            'candidates': [
                {
                    'finish_reason': GoogleFinishReason.STOP,
                    'content': {
                        'role': 'model',
                        'parts': [{'text': text} for text in texts],
                    },
                }
            ],
        }
    )