From 6f3315d74a2a9bba9a0c032947fe191d70d3a8ae Mon Sep 17 00:00:00 2001 From: pragnyanramtha Date: Sun, 17 May 2026 05:42:10 +0000 Subject: [PATCH] fix: account for operation polling timeout --- google/genai/_transformers.py | 2 +- .../tests/transformers/test_t_operation.py | 58 +++++++++++++++++++ 2 files changed, 59 insertions(+), 1 deletion(-) create mode 100644 google/genai/tests/transformers/test_t_operation.py diff --git a/google/genai/_transformers.py b/google/genai/_transformers.py index c36524100..b555e77bd 100644 --- a/google/genai/_transformers.py +++ b/google/genai/_transformers.py @@ -1186,7 +1186,7 @@ def t_resolve_operation( http_method='GET', path=name, request_dict={} ) time.sleep(delay_seconds) - total_seconds += total_seconds + total_seconds += delay_seconds # Exponential backoff delay_seconds = min( delay_seconds * LRO_POLLING_MULTIPLIER, diff --git a/google/genai/tests/transformers/test_t_operation.py b/google/genai/tests/transformers/test_t_operation.py new file mode 100644 index 000000000..687033817 --- /dev/null +++ b/google/genai/tests/transformers/test_t_operation.py @@ -0,0 +1,58 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +"""Tests for operation transformers.""" + +from unittest import mock + +import pytest + +from ... import _transformers as t + + +class _PendingOperationClient: + + def __init__(self) -> None: + self.request_count = 0 + + def request(self, **kwargs): + self.request_count += 1 + if self.request_count > 5: + raise AssertionError('operation polling did not time out') + return { + 'name': 'projects/test/locations/us-central1/operations/123', + 'done': False, + } + + +def test_resolve_operation_times_out_for_pending_operation(monkeypatch): + client = _PendingOperationClient() + + monkeypatch.setattr(t, 'LRO_POLLING_INITIAL_DELAY_SECONDS', 0.5) + monkeypatch.setattr(t, 'LRO_POLLING_MAXIMUM_DELAY_SECONDS', 0.5) + monkeypatch.setattr(t, 'LRO_POLLING_TIMEOUT_SECONDS', 1.0) + + with mock.patch.object(t.time, 'sleep') as sleep_mock: + with pytest.raises(RuntimeError, match='timed out'): + t.t_resolve_operation( + client, + { + 'name': 'projects/test/locations/us-central1/operations/123', + 'done': False, + }, + ) + + assert client.request_count == 3 + sleep_mock.assert_has_calls([mock.call(0.5), mock.call(0.5), mock.call(0.5)])