diff --git a/cirq-google/cirq_google/engine/stream_manager_test.py b/cirq-google/cirq_google/engine/stream_manager_test.py index 494b93302e7..2a440cc5724 100644 --- a/cirq-google/cirq_google/engine/stream_manager_test.py +++ b/cirq-google/cirq_google/engine/stream_manager_test.py @@ -21,6 +21,7 @@ import pytest import google.api_core.exceptions as google_exceptions +import cirq from cirq_google.engine.asyncio_executor import AsyncioExecutor from cirq_google.engine.stream_manager import ( _get_retry_request_or_raise, @@ -368,6 +369,7 @@ async def test(): ], ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) + @cirq.testing.retry_once_after_timeout def test_submit_with_retryable_stream_breakage_expects_get_result_request( self, client_constructor, error ): @@ -407,6 +409,7 @@ async def test(): ], ) @mock.patch.object(quantum, 'QuantumEngineServiceAsyncClient', autospec=True) + @cirq.testing.retry_once_after_timeout def test_submit_with_non_retryable_stream_breakage_raises_error( self, client_constructor, error ):