1212# See the License for the specific language governing permissions and
1313# limitations under the License.
1414#
15+ import asyncio
1516import cloudpickle
1617import difflib
1718import importlib
2021import sys
2122import tarfile
2223import tempfile
24+ import time
2325from typing import Any , AsyncIterable , Dict , Iterable , List , Optional
2426from unittest import mock
2527
@@ -847,6 +849,37 @@ def mock_streamer():
847849 yield stream_query_agent_engine_mock
848850
849851
852+ class _AsyncChunkIterator :
853+ """Yields response chunks asynchronously to simulate grpc.aio streams."""
854+
855+ def __init__ (self , chunks ):
856+ self ._chunks = iter (chunks )
857+
858+ def __aiter__ (self ):
859+ return self
860+
861+ async def __anext__ (self ):
862+ try :
863+ return next (self ._chunks )
864+ except StopIteration :
865+ raise StopAsyncIteration
866+
867+
868+ @pytest .fixture (scope = "function" )
869+ def async_stream_query_agent_engine_mock ():
870+ # Simulates the GAPIC async client contract of returning an awaitable that
871+ # resolves to an async iterator.
872+ async def mock_streamer (* args , ** kwargs ):
873+ return _AsyncChunkIterator (_TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE )
874+
875+ with mock .patch .object (
876+ reasoning_engine_execution_service .ReasoningEngineExecutionServiceAsyncClient ,
877+ "stream_query_reasoning_engine" ,
878+ side_effect = mock_streamer ,
879+ ) as async_stream_query_agent_engine_mock :
880+ yield async_stream_query_agent_engine_mock
881+
882+
850883@pytest .fixture (scope = "function" )
851884def get_gca_resource_mock ():
852885 with mock .patch .object (
@@ -2835,7 +2868,7 @@ async def test_async_stream_query_after_create_agent_engine_with_operation_schem
28352868 test_engine ,
28362869 test_class_method_docs ,
28372870 test_class_methods_spec ,
2838- stream_query_agent_engine_mock ,
2871+ async_stream_query_agent_engine_mock ,
28392872 ):
28402873 with mock .patch .object (
28412874 base .VertexAiResourceNoun ,
@@ -2857,7 +2890,7 @@ async def test_async_stream_query_after_create_agent_engine_with_operation_schem
28572890
28582891 assert len (results ) == 2 # Matches the length of mocked response
28592892
2860- stream_query_agent_engine_mock .assert_called_with (
2893+ async_stream_query_agent_engine_mock .assert_called_with (
28612894 request = types .StreamQueryReasoningEngineRequest (
28622895 name = _TEST_AGENT_ENGINE_RESOURCE_NAME ,
28632896 input = {"input" : _TEST_QUERY_PROMPT },
@@ -2894,7 +2927,7 @@ async def test_async_stream_query_after_update_agent_engine_with_operation_schem
28942927 test_class_methods ,
28952928 test_class_methods_spec ,
28962929 update_agent_engine_mock ,
2897- stream_query_agent_engine_mock ,
2930+ async_stream_query_agent_engine_mock ,
28982931 ):
28992932 with mock .patch .object (
29002933 base .VertexAiResourceNoun ,
@@ -2908,6 +2941,7 @@ async def test_async_stream_query_after_update_agent_engine_with_operation_schem
29082941 test_agent_engine = agent_engines .create (MethodToBeUnregisteredEngine ())
29092942 assert hasattr (test_agent_engine , _TEST_METHOD_TO_BE_UNREGISTERED_NAME )
29102943
2944+ async_client_before_update = test_agent_engine .execution_async_client
29112945 with mock .patch .object (
29122946 base .VertexAiResourceNoun ,
29132947 "_get_gca_resource" ,
@@ -2920,6 +2954,10 @@ async def test_async_stream_query_after_update_agent_engine_with_operation_schem
29202954 )
29212955 test_agent_engine .update (agent_engine = test_engine )
29222956
2957+ # Ensures update() reinitializes the async client used for streaming queries.
2958+ assert (
2959+ test_agent_engine .execution_async_client is not async_client_before_update
2960+ )
29232961 assert not hasattr (test_agent_engine , _TEST_METHOD_TO_BE_UNREGISTERED_NAME )
29242962 for method_name in test_class_methods :
29252963 invoked_method = getattr (test_agent_engine , method_name )
@@ -2929,7 +2967,7 @@ async def test_async_stream_query_after_update_agent_engine_with_operation_schem
29292967
29302968 assert len (results ) == 2 # Matches the length of mocked response
29312969
2932- stream_query_agent_engine_mock .assert_called_with (
2970+ async_stream_query_agent_engine_mock .assert_called_with (
29332971 request = types .StreamQueryReasoningEngineRequest (
29342972 name = _TEST_AGENT_ENGINE_RESOURCE_NAME ,
29352973 input = {"input" : _TEST_QUERY_PROMPT },
@@ -2965,7 +3003,7 @@ async def test_async_stream_query_agent_engine_with_operation_schema(
29653003 test_engine ,
29663004 test_class_methods ,
29673005 test_class_methods_spec ,
2968- stream_query_agent_engine_mock ,
3006+ async_stream_query_agent_engine_mock ,
29693007 ):
29703008 with mock .patch .object (
29713009 base .VertexAiResourceNoun ,
@@ -2987,14 +3025,82 @@ async def test_async_stream_query_agent_engine_with_operation_schema(
29873025
29883026 assert len (results ) == 2 # Matches the length of mocked response
29893027
2990- stream_query_agent_engine_mock .assert_called_with (
3028+ async_stream_query_agent_engine_mock .assert_called_with (
29913029 request = types .StreamQueryReasoningEngineRequest (
29923030 name = _TEST_AGENT_ENGINE_RESOURCE_NAME ,
29933031 input = {"input" : _TEST_QUERY_PROMPT },
29943032 class_method = method_name ,
29953033 )
29963034 )
29973035
3036+ @pytest .mark .asyncio
3037+ async def test_async_stream_query_keeps_event_loop_responsive (self ):
3038+ num_chunks = 5
3039+ chunk_delay_s = 0.02
3040+ watchdog_tick_s = 0.001
3041+
3042+ class _SlowAsyncStream :
3043+ """Yields chunks with non-blocking delays to simulate a grpc.aio stream."""
3044+
3045+ def __init__ (self ):
3046+ self ._remaining = num_chunks
3047+
3048+ def __aiter__ (self ):
3049+ return self
3050+
3051+ async def __anext__ (self ):
3052+ if self ._remaining == 0 :
3053+ raise StopAsyncIteration
3054+ self ._remaining -= 1
3055+ await asyncio .sleep (chunk_delay_s )
3056+ return _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE [0 ]
3057+
3058+ async def mock_async_streamer (* args , ** kwargs ):
3059+ return _SlowAsyncStream ()
3060+
3061+ def mock_blocking_sync_streamer (* args , ** kwargs ):
3062+ """Simulates a blocking sync client that blocks the thread during stream iteration."""
3063+ for _ in range (num_chunks ):
3064+ time .sleep (chunk_delay_s )
3065+ yield _TEST_AGENT_ENGINE_STREAM_QUERY_RESPONSE [0 ]
3066+
3067+ test_agent_engine = mock .MagicMock ()
3068+ test_agent_engine .resource_name = _TEST_AGENT_ENGINE_RESOURCE_NAME
3069+ test_agent_engine .execution_async_client .stream_query_reasoning_engine = (
3070+ mock_async_streamer
3071+ )
3072+ test_agent_engine .execution_api_client .stream_query_reasoning_engine = (
3073+ mock_blocking_sync_streamer
3074+ )
3075+
3076+ ticks = 0
3077+ stream_finished = asyncio .Event ()
3078+
3079+ async def watchdog ():
3080+ nonlocal ticks
3081+ while not stream_finished .is_set ():
3082+ await asyncio .sleep (watchdog_tick_s )
3083+ ticks += 1
3084+
3085+ watchdog_task = asyncio .create_task (watchdog ())
3086+ invoked_method = _agent_engines ._wrap_async_stream_query_operation (
3087+ method_name = _TEST_DEFAULT_ASYNC_STREAM_METHOD_NAME
3088+ )
3089+ results = [
3090+ chunk
3091+ async for chunk in invoked_method (
3092+ test_agent_engine , input = _TEST_QUERY_PROMPT
3093+ )
3094+ ]
3095+ stream_finished .set ()
3096+ await watchdog_task
3097+
3098+ assert len (results ) == num_chunks
3099+ # A blocking client starves the watchdog task by never yielding to the
3100+ # event loop. The threshold uses a loose lower bound to tolerate CI
3101+ # scheduling delays while verifying loop responsiveness.
3102+ assert ticks > num_chunks * 2
3103+
29983104 # pytest does not allow absl.testing.parameterized.named_parameters.
29993105 @pytest .mark .parametrize (
30003106 "test_case_name, test_engine, test_class_method_docs, test_class_methods_spec" ,
0 commit comments