Skip to content

Commit 963769b

Browse files
Adewale-1Adewale-1
committed
fix(tests): update test mocks to match api_client.request interface
Tests were mocking api_client.call_api but implementation uses api_client.request with response.body as a JSON string.
1 parent e6fe30c commit 963769b

File tree

1 file changed

+145
-189
lines changed

1 file changed

+145
-189
lines changed

‎google/genai/tests/models/test_file_state_handling.py‎

Lines changed: 145 additions & 189 deletions
Original file line numberDiff line numberDiff line change
@@ -16,207 +16,163 @@
1616

1717
"""Tests for file state handling in content generation."""
1818

19+
import json
1920
import unittest
2021
from unittest import mock
21-
import time
2222

2323
import pytest
2424

25-
from google.genai import types
2625
from google.genai import errors
26+
from google.genai import types
2727
from google.genai.models import _ensure_file_active, _process_contents_for_generation
2828
from google.genai.types import FileState
2929

3030

31+
def _make_response_body(state: str, error_message: str = None) -> bytes:
32+
"""Create a mock API response body for a file."""
33+
data = {
34+
'name': 'files/test123',
35+
'displayName': 'Test File',
36+
'mimeType': 'video/mp4',
37+
'state': state,
38+
}
39+
if error_message:
40+
data['error'] = {'message': error_message}
41+
return json.dumps(data).encode()
42+
43+
3144
class TestFileStateHandling(unittest.TestCase):
32-
"""Test file state handling functionality."""
33-
34-
def setUp(self):
35-
"""Set up test fixtures."""
36-
self.api_client = mock.MagicMock()
37-
self.file_obj = types.File(
38-
name="files/test123",
39-
display_name="Test File",
40-
mime_type="video/mp4",
41-
uri="https://example.com/files/test123",
42-
state=types.FileState.PROCESSING,
43-
)
44-
45-
@mock.patch("google.genai.files._File_from_mldev")
46-
def test_ensure_file_active_with_processing_file(self, mock_file_from_mldev):
47-
"""Test that _ensure_file_active properly handles a file in PROCESSING state."""
48-
49-
response_mock = mock.MagicMock()
50-
response_mock.json = {
51-
"file": {
52-
"name": "files/test123",
53-
"displayName": "Test File",
54-
"mimeType": "video/mp4",
55-
"state": "ACTIVE",
56-
}
57-
}
58-
self.api_client.call_api.return_value = response_mock
59-
60-
# Set up the mock to return a dict that will create an ACTIVE file
61-
mock_file_from_mldev.return_value = {
62-
"name": "files/test123",
63-
"display_name": "Test File",
64-
"mime_type": "video/mp4",
65-
"state": types.FileState.ACTIVE,
66-
}
67-
68-
result = _ensure_file_active(
69-
self.api_client, self.file_obj, max_retries=1, retry_delay_seconds=0.1
70-
)
71-
72-
self.api_client.call_api.assert_called_once_with(
73-
method="GET",
74-
url="files/test123",
75-
api_client_type="mldev",
76-
)
77-
78-
# Verify the result has ACTIVE state
79-
self.assertEqual(result.state, types.FileState.ACTIVE)
80-
81-
def test_ensure_file_active_with_failed_file(self):
82-
"""Test that _ensure_file_active properly handles a file in FAILED state."""
83-
84-
response_mock = mock.MagicMock()
85-
response_mock.json = {
86-
"file": {
87-
"name": "files/test123",
88-
"displayName": "Test File",
89-
"mimeType": "video/mp4",
90-
"state": "FAILED",
91-
"error": {"message": "File processing failed"},
92-
}
93-
}
94-
95-
# Set up a side effect for call_api that returns the response with FAILED state
96-
def mock_call_api(**kwargs):
97-
# Only return the mock for the expected file API call
98-
if kwargs.get("method") == "GET" and "files/" in kwargs.get("url", ""):
99-
return response_mock
100-
return mock.DEFAULT
101-
102-
self.api_client.call_api.side_effect = mock_call_api
103-
104-
105-
with pytest.raises(errors.FileProcessingError) as excinfo:
106-
_ensure_file_active(
107-
self.api_client, self.file_obj, max_retries=1, retry_delay_seconds=0.1
108-
)
109-
110-
assert "File processing failed" in str(excinfo.value)
111-
112-
def test_ensure_file_active_with_retries_exhausted(self):
113-
"""Test that _ensure_file_active handles a file that remains in PROCESSING state after all retries."""
114-
# Mock the response for file state check - file stays in PROCESSING
115-
response_mock = mock.MagicMock()
116-
response_mock.json = {
117-
"file": {
118-
"name": "files/test123",
119-
"displayName": "Test File",
120-
"mimeType": "video/mp4",
121-
"state": "PROCESSING",
122-
}
123-
}
124-
self.api_client.call_api.return_value = response_mock
125-
126-
# Call the function
127-
result = _ensure_file_active(
128-
self.api_client, self.file_obj, max_retries=2, retry_delay_seconds=0.1
129-
)
130-
131-
# Verify the file state was checked multiple times
132-
self.assertEqual(self.api_client.call_api.call_count, 2)
133-
134-
# Verify the original file is returned
135-
self.assertEqual(result, self.file_obj)
136-
self.assertEqual(result.state, types.FileState.PROCESSING)
137-
138-
def test_ensure_file_active_with_already_active_file(self):
139-
"""Test that _ensure_file_active returns immediately for an already ACTIVE file."""
140-
active_file = types.File(
141-
name="files/active123",
142-
display_name="Active File",
143-
mime_type="video/mp4",
144-
state=types.FileState.ACTIVE,
145-
)
146-
147-
result = _ensure_file_active(
148-
self.api_client, active_file, max_retries=1, retry_delay_seconds=0.1
149-
)
150-
151-
# Verify no API calls were made
152-
self.api_client.call_api.assert_not_called()
153-
154-
# Verify the original file is returned unchanged
155-
self.assertEqual(result, active_file)
156-
self.assertEqual(result.state, types.FileState.ACTIVE)
45+
"""Test file state handling functionality."""
46+
47+
def setUp(self):
48+
"""Set up test fixtures."""
49+
self.api_client = mock.MagicMock()
50+
self.file_obj = types.File(
51+
name='files/test123',
52+
display_name='Test File',
53+
mime_type='video/mp4',
54+
uri='https://example.com/files/test123',
55+
state=types.FileState.PROCESSING,
56+
)
57+
58+
def test_ensure_file_active_with_processing_file(self):
59+
"""Test that _ensure_file_active waits for a PROCESSING file to become ACTIVE."""
60+
response_mock = mock.MagicMock()
61+
response_mock.body = _make_response_body('ACTIVE')
62+
self.api_client.request.return_value = response_mock
63+
64+
result = _ensure_file_active(
65+
self.api_client, self.file_obj, max_retries=1, retry_delay_seconds=0
66+
)
67+
68+
self.api_client.request.assert_called_once_with(
69+
'GET', 'files/test123', {}, None
70+
)
71+
self.assertEqual(result.state, types.FileState.ACTIVE)
72+
73+
def test_ensure_file_active_with_failed_file(self):
74+
"""Test that _ensure_file_active raises FileProcessingError for a FAILED file."""
75+
response_mock = mock.MagicMock()
76+
response_mock.body = _make_response_body(
77+
'FAILED', error_message='File processing failed'
78+
)
79+
self.api_client.request.return_value = response_mock
80+
81+
with pytest.raises(errors.FileProcessingError) as excinfo:
82+
_ensure_file_active(
83+
self.api_client, self.file_obj, max_retries=1, retry_delay_seconds=0
84+
)
85+
86+
assert 'File processing failed' in str(excinfo.value)
87+
88+
def test_ensure_file_active_with_retries_exhausted(self):
89+
"""Test that _ensure_file_active returns original file after exhausting retries."""
90+
response_mock = mock.MagicMock()
91+
response_mock.body = _make_response_body('PROCESSING')
92+
self.api_client.request.return_value = response_mock
93+
94+
result = _ensure_file_active(
95+
self.api_client, self.file_obj, max_retries=2, retry_delay_seconds=0
96+
)
97+
98+
self.assertEqual(self.api_client.request.call_count, 2)
99+
self.assertEqual(result, self.file_obj)
100+
self.assertEqual(result.state, types.FileState.PROCESSING)
101+
102+
def test_ensure_file_active_with_already_active_file(self):
103+
"""Test that _ensure_file_active returns immediately for an already ACTIVE file."""
104+
active_file = types.File(
105+
name='files/active123',
106+
display_name='Active File',
107+
mime_type='video/mp4',
108+
state=types.FileState.ACTIVE,
109+
)
110+
111+
result = _ensure_file_active(
112+
self.api_client, active_file, max_retries=1, retry_delay_seconds=0
113+
)
114+
115+
self.api_client.request.assert_not_called()
116+
self.assertEqual(result, active_file)
117+
self.assertEqual(result.state, types.FileState.ACTIVE)
157118

158119

159120
class TestProcessContentsFunction(unittest.TestCase):
160-
"""Test the _process_contents_for_generation function."""
161-
162-
def setUp(self):
163-
"""Set up test fixtures."""
164-
self.api_client = mock.MagicMock()
165-
self.processing_file = types.File(
166-
name="files/processing123",
167-
display_name="Processing File",
168-
mime_type="video/mp4",
169-
uri="https://example.com/files/processing123",
170-
state=types.FileState.PROCESSING
171-
)
172-
self.active_file = types.File(
173-
name="files/active123",
174-
display_name="Active File",
175-
mime_type="video/mp4",
176-
uri="https://example.com/files/active123",
177-
state=types.FileState.ACTIVE
178-
)
179-
180-
def test_process_contents_with_files(self):
181-
"""Test that _process_contents_for_generation can handle various file scenarios."""
182-
# Scenarios:
183-
# - File directly in content list
184-
# - File in content parts
185-
# - Multiple files in different parts
186-
187-
# Test data
188-
file_in_list = [self.processing_file, "Process this file"]
189-
190-
file_in_parts = types.Content(
191-
role="user",
192-
parts=[types.Part(text="Here's a video:"), self.processing_file]
193-
)
194-
195-
multiple_files = [
196-
types.Content(
197-
role="user",
198-
parts=[types.Part(text="First video:"), self.processing_file]
121+
"""Test the _process_contents_for_generation function."""
122+
123+
def setUp(self):
124+
"""Set up test fixtures."""
125+
self.api_client = mock.MagicMock()
126+
self.processing_file = types.File(
127+
name='files/processing123',
128+
display_name='Processing File',
129+
mime_type='video/mp4',
130+
uri='https://example.com/files/processing123',
131+
state=types.FileState.PROCESSING,
132+
)
133+
self.active_file = types.File(
134+
name='files/active123',
135+
display_name='Active File',
136+
mime_type='video/mp4',
137+
uri='https://example.com/files/active123',
138+
state=types.FileState.ACTIVE,
139+
)
140+
141+
def test_process_contents_with_files(self):
142+
"""Test that _process_contents_for_generation can handle various file scenarios."""
143+
file_in_list = [self.processing_file, 'Process this file']
144+
file_in_parts = types.Content(
145+
role='user',
146+
parts=[types.Part(text="Here's a video:"), self.processing_file],
147+
)
148+
multiple_files = [
149+
types.Content(
150+
role='user',
151+
parts=[types.Part(text='First video:'), self.processing_file],
152+
),
153+
types.Content(
154+
role='user',
155+
parts=[types.Part(text='Second video:'), self.active_file],
156+
),
157+
]
158+
159+
with mock.patch(
160+
'google.genai.models._ensure_file_active', side_effect=lambda client, f: f
161+
):
162+
for test_content in [file_in_list, file_in_parts, multiple_files]:
163+
with mock.patch(
164+
'google.genai.models.t.t_contents',
165+
return_value=(
166+
test_content
167+
if isinstance(test_content, list)
168+
else [test_content]
199169
),
200-
types.Content(
201-
role="user",
202-
parts=[types.Part(text="Second video:"), self.active_file]
203-
)
204-
]
205-
206-
# Mock _ensure_file_active to return the file unchanged
207-
# This allows us to test the function without changing file states
208-
with mock.patch("google.genai.models._ensure_file_active",
209-
side_effect=lambda client, file: file):
210-
211-
# Test all three cases
212-
for test_content in [file_in_list, file_in_parts, multiple_files]:
213-
with mock.patch("google.genai.models.t.t_contents",
214-
return_value=test_content if isinstance(test_content, list) else [test_content]):
215-
# Just verify it runs without errors
216-
result = _process_contents_for_generation(self.api_client, test_content)
217-
# Basic assertion that we got something back
218-
self.assertTrue(result)
219-
220-
221-
if __name__ == "__main__":
222-
unittest.main()
170+
):
171+
result = _process_contents_for_generation(
172+
self.api_client, test_content
173+
)
174+
self.assertTrue(result)
175+
176+
177+
if __name__ == '__main__':
178+
unittest.main()

0 commit comments

Comments
 (0)