Skip to content

Commit 0d658dc

Browse files
committed
Add tests for file search edge cases to achieve 100% coverage
Add two tests to cover missing branches in file search streaming handler: - test_google_model_file_search_empty_grounding_chunks: Tests when grounding_chunks is empty/None - test_google_model_file_search_empty_retrieved_contexts: Tests when retrieved_contexts is empty (grounding_chunks exist but have no retrieved_context) These tests follow the existing pattern of mocking GenerateContentResponse for edge cases (similar to test_gemini_streamed_response_emits_text_events_for_non_empty_parts). This ensures 100% code coverage for the file search tool implementation.
1 parent 925a909 commit 0d658dc

File tree

1 file changed

+96
-0
lines changed

1 file changed

+96
-0
lines changed

tests/models/test_google.py

Lines changed: 96 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4113,6 +4113,102 @@ async def test_google_model_file_search_tool_stream(allow_model_requests: None,
41134113
await client.aio.file_search_stores.delete(name=store.name, config={'force': True})
41144114

41154115

4116+
async def test_google_model_file_search_empty_grounding_chunks():
4117+
"""Test that file search handles empty grounding_chunks gracefully."""
4118+
# Create a response with grounding_metadata but empty grounding_chunks
4119+
response = GenerateContentResponse.model_validate(
4120+
{
4121+
'response_id': 'test-123',
4122+
'model_version': 'gemini-test',
4123+
'usage_metadata': GenerateContentResponseUsageMetadata(
4124+
prompt_token_count=0,
4125+
candidates_token_count=0,
4126+
),
4127+
'candidates': [
4128+
{
4129+
'finish_reason': GoogleFinishReason.STOP,
4130+
'content': {
4131+
'role': 'model',
4132+
'parts': [{'text': 'Some text'}],
4133+
},
4134+
'grounding_metadata': {
4135+
'grounding_chunks': [], # Empty grounding_chunks
4136+
},
4137+
}
4138+
],
4139+
}
4140+
)
4141+
4142+
async def response_iterator() -> AsyncIterator[GenerateContentResponse]:
4143+
yield response
4144+
4145+
streamed_response = GeminiStreamedResponse(
4146+
model_request_parameters=ModelRequestParameters(),
4147+
_model_name='gemini-test',
4148+
_response=response_iterator(),
4149+
_timestamp=IsDatetime(),
4150+
_provider_name='test-provider',
4151+
_provider_url='',
4152+
)
4153+
# Set _file_search_tool_call_id to trigger the code path
4154+
streamed_response._file_search_tool_call_id = 'test-tool-call-id' # pyright: ignore[reportPrivateUsage]
4155+
4156+
events = [event async for event in streamed_response._get_event_iterator()] # pyright: ignore[reportPrivateUsage]
4157+
# Should not crash and should emit text event
4158+
assert len(events) > 0
4159+
assert any(isinstance(event.part, TextPart) if hasattr(event, 'part') else False for event in events)
4160+
4161+
4162+
async def test_google_model_file_search_empty_retrieved_contexts():
4163+
"""Test that file search handles empty retrieved_contexts gracefully."""
4164+
# Create a response with grounding_chunks but no retrieved_context
4165+
response = GenerateContentResponse.model_validate(
4166+
{
4167+
'response_id': 'test-123',
4168+
'model_version': 'gemini-test',
4169+
'usage_metadata': GenerateContentResponseUsageMetadata(
4170+
prompt_token_count=0,
4171+
candidates_token_count=0,
4172+
),
4173+
'candidates': [
4174+
{
4175+
'finish_reason': GoogleFinishReason.STOP,
4176+
'content': {
4177+
'role': 'model',
4178+
'parts': [{'text': 'Some text'}],
4179+
},
4180+
'grounding_metadata': {
4181+
'grounding_chunks': [
4182+
{
4183+
# Chunk without retrieved_context
4184+
}
4185+
],
4186+
},
4187+
}
4188+
],
4189+
}
4190+
)
4191+
4192+
async def response_iterator() -> AsyncIterator[GenerateContentResponse]:
4193+
yield response
4194+
4195+
streamed_response = GeminiStreamedResponse(
4196+
model_request_parameters=ModelRequestParameters(),
4197+
_model_name='gemini-test',
4198+
_response=response_iterator(),
4199+
_timestamp=IsDatetime(),
4200+
_provider_name='test-provider',
4201+
_provider_url='',
4202+
)
4203+
# Set _file_search_tool_call_id to trigger the code path
4204+
streamed_response._file_search_tool_call_id = 'test-tool-call-id' # pyright: ignore[reportPrivateUsage]
4205+
4206+
events = [event async for event in streamed_response._get_event_iterator()] # pyright: ignore[reportPrivateUsage]
4207+
# Should not crash and should emit text event
4208+
assert len(events) > 0
4209+
assert any(isinstance(event.part, TextPart) if hasattr(event, 'part') else False for event in events)
4210+
4211+
41164212
async def test_cache_point_filtering():
41174213
"""Test that CachePoint is filtered out in Google internal method."""
41184214
from pydantic_ai import CachePoint

0 commit comments

Comments
 (0)