From ca895d66d9acd36a3ab37f40f59536ebda506f10 Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 10 Sep 2026 00:29:53 -0500 Subject: [PATCH 1/4] Wait for idle instead of polling traced queries in LangSmith tests TestComprehensiveTracing and the plugin end-to-end test waited for the workflow to reach its signal wait by polling is_waiting_for_signal every second through the traced client, then asserted that every QueryWorkflow run had exactly one HandleQuery child. Temporal does not guarantee that: a query issued while a workflow task is pending or started is buffered by the server and attached to every workflow task start until it completes, including the retry after the task times out or its completion is rejected, so the worker evaluates it twice (two HandleQuery runs). If the client's 30s RPC timeout fires first, the server drops the query before it is dispatched at all (no HandleQuery run). CI logs of the failing runs show exactly this: a QueryWorkflow call lasting 29.998s followed by "Error reporting WFT to server" for the test workflow. The same stalls made the unguarded raw query fail with "Timeout expired" and, once a task reached attempt 3, made the server reject updates with "Workflow Task in failed state". Readiness is now observed through the raw client's describe(), which is not traced and involves no workflow task, and the traced query is issued exactly once after the workflow is idle, so it is dispatched directly as a query task and cannot be re-attached to a retried workflow task. Updates are only sent once the workflow is idle again after the signal. The query assertion becomes an exact hierarchy check instead of a loop over a variable number of polls. --- tests/contrib/langsmith/test_integration.py | 116 ++++++++++---------- tests/contrib/langsmith/test_plugin.py | 36 +++--- 2 files changed, 77 insertions(+), 75 deletions(-) diff --git a/tests/contrib/langsmith/test_integration.py b/tests/contrib/langsmith/test_integration.py index f48d9d6ac..196d247a1 100644 --- a/tests/contrib/langsmith/test_integration.py +++ b/tests/contrib/langsmith/test_integration.py @@ -16,13 +16,12 @@ from temporalio import activity, common, nexus, workflow from temporalio.client import ( Client, + WorkflowExecutionStatus, WorkflowFailureError, WorkflowHandle, - WorkflowQueryFailedError, ) from temporalio.contrib.langsmith import LangSmithPlugin from temporalio.exceptions import ApplicationError -from temporalio.service import RPCError from temporalio.testing import WorkflowEnvironment from tests.contrib.langsmith.conftest import ( InMemoryRunCollector, @@ -312,22 +311,28 @@ def _make_temporal_client( return Client(**config) -@traceable(name="poll_query") -async def _poll_query( - handle: WorkflowHandle[Any, Any], - query: Callable[..., Any], - *, - expected: Any = True, -) -> bool: - """Poll a workflow query until it returns the expected value.""" +async def _wait_for_workflow_idle(handle: WorkflowHandle[Any, Any]) -> None: + """Wait until the running workflow has no pending task, activity, child or Nexus op.""" while True: - try: - result = await handle.query(query) - if result == expected: - return True - except (WorkflowQueryFailedError, RPCError): - pass # Query not yet available (workflow hasn't started) - await asyncio.sleep(1) + desc = await handle.describe() + assert desc.status == WorkflowExecutionStatus.RUNNING, desc.status + raw = desc.raw_description + if not ( + raw.HasField("pending_workflow_task") + or raw.pending_activities + or raw.pending_children + or raw.pending_nexus_operations + ): + return + await asyncio.sleep(0.2) + + +@traceable(name="query_pipeline") +async def _query_pipeline( + handle: WorkflowHandle[Any, Any], query: Callable[..., Any] +) -> Any: + """Run a single workflow query under a @traceable root.""" + return await handle.query(query) # --------------------------------------------------------------------------- @@ -563,7 +568,7 @@ async def test_comprehensive_with_temporal_runs( """Full trace hierarchy with worker restart mid-workflow. user_pipeline only wraps start_workflow (completing before the worker - starts), so poll/signal/query traces are naturally separate root traces. + starts), so query/signal/update traces are naturally separate root traces. """ if env.supports_time_skipping: pytest.skip("Time-skipping server doesn't persist headers.") @@ -575,6 +580,8 @@ async def test_comprehensive_with_temporal_runs( temporal_client_1 = _make_temporal_client( client, mock_ls, add_temporal_runs=True ) + # Raw-client handle (no LangSmith interceptor) for untraced readiness checks + raw_handle = client.get_workflow_handle(workflow_id) @traceable(name="user_pipeline") async def user_pipeline() -> WorkflowHandle[Any, Any]: @@ -588,7 +595,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: # Start workflow — no worker yet, just a server RPC handle = await user_pipeline() - # Phase 1: worker picks up workflow, poll until signal wait + # Phase 1: worker runs the workflow up to the signal wait async with new_worker( temporal_client_1, ComprehensiveWorkflow, @@ -602,14 +609,14 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: make_nexus_endpoint_name(worker.task_queue), worker.task_queue, ) - assert await _poll_query( - handle, - ComprehensiveWorkflow.is_waiting_for_signal, - expected=True, + await _wait_for_workflow_idle(raw_handle) + assert await _query_pipeline( + handle, ComprehensiveWorkflow.is_waiting_for_signal ), "Workflow never reached signal wait point" - # Raw-client query (no LangSmith interceptor) — root-level trace - raw_handle = client.get_workflow_handle(workflow_id) - await raw_handle.query(ComprehensiveWorkflow.is_waiting_for_signal) + # Raw-client query — root-level trace + assert await raw_handle.query( + ComprehensiveWorkflow.is_waiting_for_signal + ) # Phase 2: fresh worker, signal to resume, complete temporal_client_2 = _make_temporal_client( @@ -627,6 +634,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: handle_2 = temporal_client_2.get_workflow_handle(workflow_id) await handle_2.query(ComprehensiveWorkflow.my_query) await handle_2.signal(ComprehensiveWorkflow.my_signal, "hello") + await _wait_for_workflow_idle(raw_handle) await handle_2.execute_update( ComprehensiveWorkflow.my_unvalidated_update, "test" ) @@ -704,18 +712,15 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: ] assert_trace_hierarchy(workflow_trace_trees, expected_workflow) - # poll_query trace (separate root, variable number of iterations) - poll_trace_trees = find_trace_trees(trace_trees, "poll_query") - assert len(poll_trace_trees) == 1 - poll = poll_trace_trees[0] - assert poll.name == "poll_query" - poll_children = poll.children - for poll_child in poll_children: - assert poll_child.name == "QueryWorkflow:is_waiting_for_signal" - assert [child.name for child in poll_child.children] == [ - "HandleQuery:is_waiting_for_signal" - ] - assert not poll_child.children[0].children + # query_pipeline trace: the worker-side handler nests under the client query + assert_trace_hierarchy( + find_trace_trees(trace_trees, "query_pipeline"), + [ + "query_pipeline", + " QueryWorkflow:is_waiting_for_signal", + " HandleQuery:is_waiting_for_signal", + ], + ) # Raw-client query — no parent context, appears as root raw_query_trace_trees = [ @@ -777,7 +782,7 @@ async def test_comprehensive_without_temporal_runs( """Same workflow with add_temporal_runs=False and worker restart. Only @traceable runs appear. Context propagation via headers still works. - user_pipeline only wraps start_workflow, so poll traces are separate roots. + user_pipeline only wraps start_workflow, so the query trace is a separate root. """ if env.supports_time_skipping: pytest.skip("Time-skipping server doesn't persist headers.") @@ -789,6 +794,8 @@ async def test_comprehensive_without_temporal_runs( temporal_client_1 = _make_temporal_client( client, mock_ls, add_temporal_runs=False ) + # Raw-client handle (no LangSmith interceptor) for untraced readiness checks + raw_handle = client.get_workflow_handle(workflow_id) @traceable(name="user_pipeline") async def user_pipeline() -> WorkflowHandle[Any, Any]: @@ -801,7 +808,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: with tracing_context(client=mock_ls, enabled=True): handle = await user_pipeline() - # Phase 1: worker picks up workflow, poll until signal wait + # Phase 1: worker runs the workflow up to the signal wait async with new_worker( temporal_client_1, ComprehensiveWorkflow, @@ -815,13 +822,13 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: make_nexus_endpoint_name(worker.task_queue), worker.task_queue, ) + await _wait_for_workflow_idle(raw_handle) # Raw-client query — no interceptor, produces nothing - raw_handle = client.get_workflow_handle(workflow_id) - await raw_handle.query(ComprehensiveWorkflow.is_waiting_for_signal) - assert await _poll_query( - handle, - ComprehensiveWorkflow.is_waiting_for_signal, - expected=True, + assert await raw_handle.query( + ComprehensiveWorkflow.is_waiting_for_signal + ) + assert await _query_pipeline( + handle, ComprehensiveWorkflow.is_waiting_for_signal ), "Workflow never reached signal wait point" # Phase 2: fresh worker, signal to resume, complete @@ -839,6 +846,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: ): handle_2 = temporal_client_2.get_workflow_handle(workflow_id) await handle_2.signal(ComprehensiveWorkflow.my_signal, "hello") + await _wait_for_workflow_idle(raw_handle) await handle_2.execute_update( ComprehensiveWorkflow.my_unvalidated_update, "test" ) @@ -883,10 +891,10 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: ] assert_trace_hierarchy(workflow_trace_trees, expected_workflow) - # Poll query — separate root, just the @traceable wrapper, no Temporal children - poll_trace_trees = find_trace_trees(trace_trees, "poll_query") - assert len(poll_trace_trees) == 1 - assert_trace_hierarchy(poll_trace_trees, ["poll_query"]) + # Query — separate root, just the @traceable wrapper, no Temporal children + assert_trace_hierarchy( + find_trace_trees(trace_trees, "query_pipeline"), ["query_pipeline"] + ) # --------------------------------------------------------------------------- @@ -1252,13 +1260,7 @@ async def test_temporal_prefixed_query_not_traced( task_queue=worker.task_queue, ) - # Wait for workflow to start by polling the user query - assert await _poll_query( - handle, - QueryFilteringWorkflow.my_query, - expected="query-result", - ), "Workflow never started" - + await _wait_for_workflow_idle(handle) collector.clear() # Built-in queries — should NOT be traced diff --git a/tests/contrib/langsmith/test_plugin.py b/tests/contrib/langsmith/test_plugin.py index 6e3cb2e86..b1f1e7593 100644 --- a/tests/contrib/langsmith/test_plugin.py +++ b/tests/contrib/langsmith/test_plugin.py @@ -21,7 +21,8 @@ NexusService, TraceableActivityWorkflow, _make_client_and_collector, - _poll_query, + _query_pipeline, + _wait_for_workflow_idle, nested_traceable_activity, traceable_activity, ) @@ -63,7 +64,7 @@ async def test_comprehensive_plugin_trace_hierarchy( ) -> None: """Plugin wired to a real Temporal worker produces the full trace hierarchy. - user_pipeline only wraps start_workflow, so poll/query/signal/update + user_pipeline only wraps start_workflow, so query/signal/update traces are naturally separate root traces. """ if env.supports_time_skipping: @@ -100,13 +101,15 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: make_nexus_endpoint_name(worker.task_queue), worker.task_queue, ) - assert await _poll_query( - handle, - ComprehensiveWorkflow.is_waiting_for_signal, - expected=True, + # Raw-client handle (no LangSmith interceptor) for untraced readiness checks + raw_handle = client.get_workflow_handle(workflow_id) + await _wait_for_workflow_idle(raw_handle) + assert await _query_pipeline( + handle, ComprehensiveWorkflow.is_waiting_for_signal ), "Workflow never reached signal wait point" await handle.query(ComprehensiveWorkflow.my_query) await handle.signal(ComprehensiveWorkflow.my_signal, "hello") + await _wait_for_workflow_idle(raw_handle) await handle.execute_update( ComprehensiveWorkflow.my_unvalidated_update, "test" ) @@ -184,18 +187,15 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: ] assert_trace_hierarchy(workflow_trace_trees, expected_workflow) - # poll_query trace (separate root, variable number of iterations) - poll_trace_trees = find_trace_trees(trace_trees, "poll_query") - assert len(poll_trace_trees) == 1 - poll = poll_trace_trees[0] - assert poll.name == "poll_query" - poll_children = poll.children - for poll_child in poll_children: - assert poll_child.name == "QueryWorkflow:is_waiting_for_signal" - assert [child.name for child in poll_child.children] == [ - "HandleQuery:is_waiting_for_signal" - ] - assert not poll_child.children[0].children + # query_pipeline trace: the worker-side handler nests under the client query + assert_trace_hierarchy( + find_trace_trees(trace_trees, "query_pipeline"), + [ + "query_pipeline", + " QueryWorkflow:is_waiting_for_signal", + " HandleQuery:is_waiting_for_signal", + ], + ) # Each remaining operation is its own root trace query_trace_trees = find_trace_trees(trace_trees, "QueryWorkflow:my_query") From ac107b45ecc581ae4635fe7c28ad68ef8735cfcc Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 10 Sep 2026 02:19:42 -0500 Subject: [PATCH 2/4] Expect the RunWorkflow root in the built-in query filtering test The worker's RunWorkflow run reaches the collector asynchronously and can arrive after describe() already reports the first task complete, so clearing the collector after the idle wait raced it and left a stray root on fast runners. Nothing needs clearing now that readiness no longer polls a traced query, so assert the full hierarchy instead. --- tests/contrib/langsmith/test_integration.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/contrib/langsmith/test_integration.py b/tests/contrib/langsmith/test_integration.py index 196d247a1..669b9f931 100644 --- a/tests/contrib/langsmith/test_integration.py +++ b/tests/contrib/langsmith/test_integration.py @@ -1261,7 +1261,6 @@ async def test_temporal_prefixed_query_not_traced( ) await _wait_for_workflow_idle(handle) - collector.clear() # Built-in queries — should NOT be traced await handle.query("__temporal_workflow_metadata") @@ -1272,10 +1271,11 @@ async def test_temporal_prefixed_query_not_traced( await handle.signal(QueryFilteringWorkflow.complete) assert await handle.result() == "done" - # Built-in queries should be absent; only user query and signal remain. + # The built-in query leaves no run; everything else the worker did is here. assert_trace_hierarchy( build_trace_trees(collector), [ + "RunWorkflow:QueryFilteringWorkflow", "HandleQuery:my_query", "HandleSignal:complete", ], From f587719a4c8e86b30176687bad7173f4edae94d3 Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 10 Sep 2026 02:59:52 -0500 Subject: [PATCH 3/4] Stop activity retries from adding runs to the LangSmith trace assertions The traced activities finish instantly; their 10s start-to-close bound only satisfies Temporal. On a stalled runner an attempt timed out, the retry ran the activity again, and the exact hierarchy assertion saw one extra run. Give the bound a minute and disable retries so a failed attempt fails loudly. --- tests/contrib/langsmith/test_integration.py | 25 +++++++++++++-------- 1 file changed, 16 insertions(+), 9 deletions(-) diff --git a/tests/contrib/langsmith/test_integration.py b/tests/contrib/langsmith/test_integration.py index 669b9f931..01e047e2b 100644 --- a/tests/contrib/langsmith/test_integration.py +++ b/tests/contrib/langsmith/test_integration.py @@ -77,7 +77,8 @@ class TraceableActivityWorkflow: async def run(self, _input: str = "") -> str: return await workflow.execute_activity( traceable_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), + retry_policy=common.RetryPolicy(maximum_attempts=1), ) @@ -111,7 +112,8 @@ class SimpleWorkflow: async def run(self) -> str: result = await workflow.execute_activity( simple_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), + retry_policy=common.RetryPolicy(maximum_attempts=1), ) return result @@ -126,7 +128,8 @@ async def _step_with_activity() -> str: """A @traceable step that wraps an activity call.""" return await workflow.execute_activity( nested_traceable_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), + retry_policy=common.RetryPolicy(maximum_attempts=1), ) @@ -164,12 +167,14 @@ def __init__(self) -> None: async def run(self) -> str: await workflow.execute_activity( nested_traceable_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), + retry_policy=common.RetryPolicy(maximum_attempts=1), ) await _step_with_activity() await workflow.execute_local_activity( nested_traceable_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), + retry_policy=common.RetryPolicy(maximum_attempts=1), ) await _outer_chain("from-workflow") await workflow.execute_child_workflow( @@ -192,7 +197,8 @@ async def run(self) -> str: await workflow.wait_condition(lambda: self._signal_received) await workflow.execute_activity( nested_traceable_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), + retry_policy=common.RetryPolicy(maximum_attempts=1), ) await workflow.wait_condition(lambda: self._complete) return "comprehensive-done" @@ -260,7 +266,7 @@ class ActivityFailureWorkflow: async def run(self) -> str: return await workflow.execute_activity( failing_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), retry_policy=common.RetryPolicy(maximum_attempts=1), ) @@ -271,7 +277,7 @@ class BenignErrorWorkflow: async def run(self) -> str: return await workflow.execute_activity( benign_failing_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), retry_policy=common.RetryPolicy(maximum_attempts=1), ) @@ -938,7 +944,8 @@ async def run(self) -> str: # Activity with nested @traceable await workflow.execute_activity( nested_traceable_activity, - start_to_close_timeout=timedelta(seconds=10), + start_to_close_timeout=timedelta(minutes=1), + retry_policy=common.RetryPolicy(maximum_attempts=1), ) return f"{r1}|{r2}|{r3}" From 897d94827438d22f2d2c9ca0c4adad304c051886 Mon Sep 17 00:00:00 2001 From: DABH Date: Fri, 11 Sep 2026 09:24:20 -0500 Subject: [PATCH 4/4] Use the shared wait_for_workflow_idle helper Drop the private idle wait now that tests.helpers has one with a timeout and a fail-fast on a workflow that is no longer running. Give it 30s here because CI has stalled workflow tasks for tens of seconds. --- tests/contrib/langsmith/test_integration.py | 31 ++++++--------------- tests/contrib/langsmith/test_plugin.py | 8 +++--- 2 files changed, 13 insertions(+), 26 deletions(-) diff --git a/tests/contrib/langsmith/test_integration.py b/tests/contrib/langsmith/test_integration.py index 01e047e2b..d5496d7dd 100644 --- a/tests/contrib/langsmith/test_integration.py +++ b/tests/contrib/langsmith/test_integration.py @@ -2,7 +2,6 @@ from __future__ import annotations -import asyncio import uuid from collections.abc import Callable from datetime import timedelta @@ -16,7 +15,6 @@ from temporalio import activity, common, nexus, workflow from temporalio.client import ( Client, - WorkflowExecutionStatus, WorkflowFailureError, WorkflowHandle, ) @@ -29,7 +27,7 @@ find_trace_trees, make_mock_ls_client, ) -from tests.helpers import new_worker +from tests.helpers import new_worker, wait_for_workflow_idle from tests.helpers.nexus import make_nexus_endpoint_name from tests.helpers.trace import assert_trace_hierarchy @@ -317,20 +315,9 @@ def _make_temporal_client( return Client(**config) -async def _wait_for_workflow_idle(handle: WorkflowHandle[Any, Any]) -> None: - """Wait until the running workflow has no pending task, activity, child or Nexus op.""" - while True: - desc = await handle.describe() - assert desc.status == WorkflowExecutionStatus.RUNNING, desc.status - raw = desc.raw_description - if not ( - raw.HasField("pending_workflow_task") - or raw.pending_activities - or raw.pending_children - or raw.pending_nexus_operations - ): - return - await asyncio.sleep(0.2) +# CI has stalled workflow tasks for tens of seconds; give the untraced idle probe +# most of the 60s per-test budget instead of the helper's 10s default. +_IDLE_TIMEOUT = timedelta(seconds=30) @traceable(name="query_pipeline") @@ -615,7 +602,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: make_nexus_endpoint_name(worker.task_queue), worker.task_queue, ) - await _wait_for_workflow_idle(raw_handle) + await wait_for_workflow_idle(raw_handle, timeout=_IDLE_TIMEOUT) assert await _query_pipeline( handle, ComprehensiveWorkflow.is_waiting_for_signal ), "Workflow never reached signal wait point" @@ -640,7 +627,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: handle_2 = temporal_client_2.get_workflow_handle(workflow_id) await handle_2.query(ComprehensiveWorkflow.my_query) await handle_2.signal(ComprehensiveWorkflow.my_signal, "hello") - await _wait_for_workflow_idle(raw_handle) + await wait_for_workflow_idle(raw_handle, timeout=_IDLE_TIMEOUT) await handle_2.execute_update( ComprehensiveWorkflow.my_unvalidated_update, "test" ) @@ -828,7 +815,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: make_nexus_endpoint_name(worker.task_queue), worker.task_queue, ) - await _wait_for_workflow_idle(raw_handle) + await wait_for_workflow_idle(raw_handle, timeout=_IDLE_TIMEOUT) # Raw-client query — no interceptor, produces nothing assert await raw_handle.query( ComprehensiveWorkflow.is_waiting_for_signal @@ -852,7 +839,7 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: ): handle_2 = temporal_client_2.get_workflow_handle(workflow_id) await handle_2.signal(ComprehensiveWorkflow.my_signal, "hello") - await _wait_for_workflow_idle(raw_handle) + await wait_for_workflow_idle(raw_handle, timeout=_IDLE_TIMEOUT) await handle_2.execute_update( ComprehensiveWorkflow.my_unvalidated_update, "test" ) @@ -1267,7 +1254,7 @@ async def test_temporal_prefixed_query_not_traced( task_queue=worker.task_queue, ) - await _wait_for_workflow_idle(handle) + await wait_for_workflow_idle(handle, timeout=_IDLE_TIMEOUT) # Built-in queries — should NOT be traced await handle.query("__temporal_workflow_metadata") diff --git a/tests/contrib/langsmith/test_plugin.py b/tests/contrib/langsmith/test_plugin.py index b1f1e7593..f6f1a6186 100644 --- a/tests/contrib/langsmith/test_plugin.py +++ b/tests/contrib/langsmith/test_plugin.py @@ -17,16 +17,16 @@ find_trace_trees, ) from tests.contrib.langsmith.test_integration import ( + _IDLE_TIMEOUT, ComprehensiveWorkflow, NexusService, TraceableActivityWorkflow, _make_client_and_collector, _query_pipeline, - _wait_for_workflow_idle, nested_traceable_activity, traceable_activity, ) -from tests.helpers import new_worker +from tests.helpers import new_worker, wait_for_workflow_idle from tests.helpers.nexus import make_nexus_endpoint_name from tests.helpers.trace import assert_trace_hierarchy @@ -103,13 +103,13 @@ async def user_pipeline() -> WorkflowHandle[Any, Any]: ) # Raw-client handle (no LangSmith interceptor) for untraced readiness checks raw_handle = client.get_workflow_handle(workflow_id) - await _wait_for_workflow_idle(raw_handle) + await wait_for_workflow_idle(raw_handle, timeout=_IDLE_TIMEOUT) assert await _query_pipeline( handle, ComprehensiveWorkflow.is_waiting_for_signal ), "Workflow never reached signal wait point" await handle.query(ComprehensiveWorkflow.my_query) await handle.signal(ComprehensiveWorkflow.my_signal, "hello") - await _wait_for_workflow_idle(raw_handle) + await wait_for_workflow_idle(raw_handle, timeout=_IDLE_TIMEOUT) await handle.execute_update( ComprehensiveWorkflow.my_unvalidated_update, "test" )