CI - fix flakiness of stream_manager_test on Windows (#6749)

Some tests in cirq_google.engine.stream_manager_test show
non-reproducible TimeoutError failures on CI runners, which
typically go away on retry.  Here we retry them by default.

Fixes #6741

cirq-core/cirq/testing/__init__.py

@@ -115,7 +115,8 @@
 
 from cirq.testing.order_tester import OrderTester as OrderTester
 
-from cirq.testing.pytest_randomly_utils import (
+from cirq.testing.pytest_utils import (
+    retry_once_after_timeout as retry_once_after_timeout,
     retry_once_with_later_random_values as retry_once_with_later_random_values,
 )
 

cirq-core/cirq/testing/pytest_randomly_utils.py → cirq-core/cirq/testing/pytest_utils.py

@@ -19,6 +19,25 @@
 from typing import Any, Callable
 
 
+def retry_once_after_timeout(testfunc: Callable) -> Callable:
+    """Marks a test function for one retry if it fails with TimeoutError.
+
+    This decorator is intended for test functions which occasionally fail
+    with TimeoutError.
+    """
+
+    @functools.wraps(testfunc)
+    def wrapped_func(*args, **kwargs) -> Any:
+        try:
+            return testfunc(*args, **kwargs)
+        except TimeoutError:
+            pass
+        warnings.warn("Retrying in case we have a transitive TimeoutError")
+        return testfunc(*args, **kwargs)
+
+    return wrapped_func
+
+
 def retry_once_with_later_random_values(testfunc: Callable) -> Callable:
     """Marks a test function for one retry with later random values.
 

cirq-core/cirq/testing/pytest_randomly_utils_test.py → cirq-core/cirq/testing/pytest_utils_test.py

@@ -19,6 +19,14 @@
 import cirq
 
 
+def test_retry_once_after_timeout():
+    testfunc = Mock(side_effect=[TimeoutError("first call fails"), None])
+    decoratedfunc = cirq.testing.retry_once_after_timeout(testfunc)
+    with pytest.warns(UserWarning, match="Retrying.*transitive TimeoutError"):
+        decoratedfunc()
+    assert testfunc.call_count == 2
+
+
 def test_retry_once_with_later_random_values():
     testfunc = Mock(side_effect=[AssertionError("first call fails"), None])
     decoratedfunc = cirq.testing.retry_once_with_later_random_values(testfunc)

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
     ):