Skip to content

Commit 66ca723

Browse files
fix: account for operation polling timeout
1 parent 2910346 commit 66ca723

2 files changed

Lines changed: 59 additions & 1 deletion

File tree

google/genai/_transformers.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1186,7 +1186,7 @@ def t_resolve_operation(
11861186
http_method='GET', path=name, request_dict={}
11871187
)
11881188
time.sleep(delay_seconds)
1189-
total_seconds += total_seconds
1189+
total_seconds += delay_seconds
11901190
# Exponential backoff
11911191
delay_seconds = min(
11921192
delay_seconds * LRO_POLLING_MULTIPLIER,
Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
# Copyright 2026 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
#
15+
16+
"""Tests for operation transformers."""
17+
18+
from unittest import mock
19+
20+
import pytest
21+
22+
from ... import _transformers as t
23+
24+
25+
class _PendingOperationClient:
26+
27+
def __init__(self) -> None:
28+
self.request_count = 0
29+
30+
def request(self, **kwargs):
31+
self.request_count += 1
32+
if self.request_count > 5:
33+
raise AssertionError('operation polling did not time out')
34+
return {
35+
'name': 'projects/test/locations/us-central1/operations/123',
36+
'done': False,
37+
}
38+
39+
40+
def test_resolve_operation_times_out_for_pending_operation(monkeypatch):
41+
client = _PendingOperationClient()
42+
43+
monkeypatch.setattr(t, 'LRO_POLLING_INITIAL_DELAY_SECONDS', 0.5)
44+
monkeypatch.setattr(t, 'LRO_POLLING_MAXIMUM_DELAY_SECONDS', 0.5)
45+
monkeypatch.setattr(t, 'LRO_POLLING_TIMEOUT_SECONDS', 1.0)
46+
47+
with mock.patch.object(t.time, 'sleep') as sleep_mock:
48+
with pytest.raises(RuntimeError, match='timed out'):
49+
t.t_resolve_operation(
50+
client,
51+
{
52+
'name': 'projects/test/locations/us-central1/operations/123',
53+
'done': False,
54+
},
55+
)
56+
57+
assert client.request_count == 3
58+
sleep_mock.assert_has_calls([mock.call(0.5), mock.call(0.5), mock.call(0.5)])

0 commit comments

Comments
 (0)