diff --git a/tests/contrib/langgraph/test_summary_fn.py b/tests/contrib/langgraph/test_summary_fn.py index 687a68edb..5236dde44 100644 --- a/tests/contrib/langgraph/test_summary_fn.py +++ b/tests/contrib/langgraph/test_summary_fn.py @@ -19,7 +19,7 @@ from temporalio.contrib.langgraph import LangGraphPlugin, graph from temporalio.testing import WorkflowEnvironment from temporalio.worker import Replayer, Worker -from tests.helpers import assert_eq_eventually +from tests.helpers import wait_for_workflow_idle SummaryFn = Callable[[tuple[Any, ...], dict[str, Any]], "str | None"] @@ -235,12 +235,10 @@ class WorkflowNodeSummaryWorkflow: def __init__(self) -> None: self.app = graph("wf-node-graph").compile() self._done = False - self._invoked = False @workflow.run async def run(self, input: str) -> Any: result = await self.app.ainvoke({"value": input}) - self._invoked = True await workflow.wait_condition(lambda: self._done) return result @@ -248,14 +246,6 @@ async def run(self, input: str) -> Any: def finish(self) -> None: self._done = True - @workflow.query - def ran(self) -> bool: - return workflow.get_current_details() != "" - - @workflow.query - def invoked(self) -> bool: - return self._invoked - async def test_workflow_node_sets_current_details( client: Client, env: WorkflowEnvironment @@ -285,9 +275,7 @@ async def test_workflow_node_sets_current_details( id=f"wf-node-{uuid.uuid4()}", task_queue=task_queue, ) - await assert_eq_eventually( - True, lambda: handle.query(WorkflowNodeSummaryWorkflow.ran) - ) + await wait_for_workflow_idle(handle) md: temporalio.api.sdk.v1.WorkflowMetadata = await handle.query( "__temporal_workflow_metadata", result_type=temporalio.api.sdk.v1.WorkflowMetadata, @@ -328,9 +316,7 @@ async def test_workflow_node_clears_current_details_on_empty( id=f"wf-node-clear-{uuid.uuid4()}", task_queue=task_queue, ) - await assert_eq_eventually( - True, lambda: handle.query(WorkflowNodeSummaryWorkflow.invoked) - ) + await wait_for_workflow_idle(handle) md: temporalio.api.sdk.v1.WorkflowMetadata = await handle.query( "__temporal_workflow_metadata", result_type=temporalio.api.sdk.v1.WorkflowMetadata, diff --git a/tests/contrib/openai_agents/test_openai_tracing.py b/tests/contrib/openai_agents/test_openai_tracing.py index 28b804cc1..42cb5293b 100644 --- a/tests/contrib/openai_agents/test_openai_tracing.py +++ b/tests/contrib/openai_agents/test_openai_tracing.py @@ -10,6 +10,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from temporalio import activity, workflow +from temporalio.api.enums.v1 import EventType from temporalio.client import Client from temporalio.contrib.openai_agents import _temporal_openai_agents from temporalio.contrib.openai_agents.testing import ( @@ -24,7 +25,7 @@ ResearchWorkflow, research_mock_model, ) -from tests.helpers import assert_eq_eventually, new_worker +from tests.helpers import assert_event_subsequence, new_worker class MemoryTracingProcessor(TracingProcessor): @@ -242,11 +243,26 @@ async def simple_no_context_activity() -> str: return "success" +async def wait_for_activity_processed(client: Client, workflow_id: str) -> None: + """Wait, via an untraced history poll, until the workflow task that handled the activity result completed. + + Assumes the workflow's first completed activity is the one right before its park point, which + holds for every workflow in this module. + """ + await assert_event_subsequence( + client.get_workflow_handle(workflow_id), + [ + EventType.EVENT_TYPE_ACTIVITY_TASK_COMPLETED, + EventType.EVENT_TYPE_WORKFLOW_TASK_COMPLETED, + ], + timeout=timedelta(seconds=10), + ) + + @workflow.defn class TraceWorkflow: def __init__(self) -> None: self._proceed = False - self._ready = False @workflow.run async def run(self): @@ -256,14 +272,9 @@ async def run(self): simple_no_context_activity, start_to_close_timeout=timedelta(seconds=10), ) - self._ready = True await workflow.wait_condition(lambda: self._proceed) return "done" - @workflow.query - def ready(self) -> bool: - return self._ready - @workflow.signal def proceed(self) -> None: self._proceed = True @@ -273,7 +284,6 @@ def proceed(self) -> None: class SelfTracingWorkflow: def __init__(self) -> None: self._proceed = False - self._ready = False @workflow.run async def run(self): @@ -284,14 +294,9 @@ async def run(self): simple_no_context_activity, start_to_close_timeout=timedelta(seconds=10), ) - self._ready = True await workflow.wait_condition(lambda: self._proceed) return "done" - @workflow.query - def ready(self) -> bool: - return self._ready - @workflow.signal def proceed(self) -> None: self._proceed = True @@ -366,11 +371,7 @@ async def test_external_trace_to_workflow_spans( max_cached_workflows=0, task_queue=task_queue, ): - # Wait for workflow to be ready - async def ready() -> bool: - return await workflow_handle.query(TraceWorkflow.ready) - - await assert_eq_eventually(True, ready) + await wait_for_activity_processed(client, workflow_handle.id) # Second worker: Complete the workflow with fresh objects (new instrumentation) async with AgentEnvironment( @@ -458,11 +459,7 @@ async def test_external_trace_and_span_to_workflow_spans( max_cached_workflows=0, task_queue=task_queue, ): - # Wait for workflow to be ready - async def ready() -> bool: - return await workflow_handle.query(TraceWorkflow.ready) - - await assert_eq_eventually(True, ready) + await wait_for_activity_processed(client, workflow_handle.id) # Second worker: Complete the workflow with fresh objects (new instrumentation) async with AgentEnvironment( @@ -554,11 +551,7 @@ async def test_workflow_only_trace_to_spans( ) workflow_id = workflow_handle.id - # Wait for workflow to be ready - async def ready() -> bool: - return await workflow_handle.query(SelfTracingWorkflow.ready) - - await assert_eq_eventually(True, ready) + await wait_for_activity_processed(client, workflow_handle.id) # Second worker: Complete the workflow with fresh objects (new instrumentation) async with AgentEnvironment( @@ -805,7 +798,6 @@ def is_descendant_of(child: ReadableSpan, ancestor_span_id: int) -> bool: class OtelSpanWorkflow: def __init__(self) -> None: self._proceed = False - self._ready = False @workflow.run async def run(self): @@ -818,14 +810,9 @@ async def run(self): simple_no_context_activity, start_to_close_timeout=timedelta(seconds=10), ) - self._ready = True await workflow.wait_condition(lambda: self._proceed) return "done" - @workflow.query - def ready(self) -> bool: - return self._ready - @workflow.signal def proceed(self) -> None: self._proceed = True @@ -868,11 +855,7 @@ async def test_sdk_trace_to_otel_span_parenting( ) workflow_id = workflow_handle.id - # Wait for workflow to be ready - async def ready() -> bool: - return await workflow_handle.query(OtelSpanWorkflow.ready) - - await assert_eq_eventually(True, ready) + await wait_for_activity_processed(client, workflow_handle.id) # Second worker: Complete the workflow with fresh objects (new instrumentation) async with AgentEnvironment( diff --git a/tests/helpers/__init__.py b/tests/helpers/__init__.py index fe37296e9..a706eb5e5 100644 --- a/tests/helpers/__init__.py +++ b/tests/helpers/__init__.py @@ -16,6 +16,8 @@ cast, ) +import pytest + from temporalio.api.common.v1 import WorkflowExecution from temporalio.api.enums.v1 import EventType as EventType from temporalio.api.enums.v1 import IndexedValueType @@ -31,7 +33,12 @@ PollWorkflowExecutionUpdateRequest, UnpauseActivityRequest, ) -from temporalio.client import BuildIdOpAddNewDefault, Client, WorkflowHandle +from temporalio.client import ( + BuildIdOpAddNewDefault, + Client, + WorkflowExecutionStatus, + WorkflowHandle, +) from temporalio.common import SearchAttributeKey from temporalio.converter import DataConverter from temporalio.service import RPCError, RPCStatusCode @@ -83,9 +90,9 @@ async def assert_eventually( if timedelta(seconds=time.monotonic() - start_sec) >= timeout: raise except RPCError as e: - if retry_on_rpc_cancelled and e.status == RPCStatusCode.CANCELLED: - continue - else: + if not (retry_on_rpc_cancelled and e.status == RPCStatusCode.CANCELLED): + raise + if timedelta(seconds=time.monotonic() - start_sec) >= timeout: raise await asyncio.sleep(interval.total_seconds()) @@ -103,6 +110,44 @@ async def check() -> None: await assert_eventually(check, timeout=timeout, interval=interval) +async def wait_for_workflow_idle( + handle: WorkflowHandle[Any, Any], + *, + timeout: timedelta = timedelta(seconds=10), + interval: timedelta = timedelta(milliseconds=200), +) -> None: + """Wait until the running workflow has no pending workflow task, activity, child or Nexus op. + + Use as an untraced readiness probe before a traced query, signal or update: pass a handle + from a client without tracing interceptors. Fails fast if the workflow is no longer running. + + Limits: a workflow blocked only on a timer also counts as idle, and the time-skipping test + server never reports ``pending_workflow_task``, so this returns immediately there. + """ + + async def check() -> None: + desc = await handle.describe() + if desc.status != WorkflowExecutionStatus.RUNNING: + status = desc.status.name if desc.status is not None else "UNKNOWN" + pytest.fail(f"Workflow {handle.id} is {status}, not RUNNING") + raw = desc.raw_description + pending = [ + name + for name, present in ( + ("workflow task", raw.HasField("pending_workflow_task")), + ("activities", bool(raw.pending_activities)), + ("child workflows", bool(raw.pending_children)), + ("Nexus operations", bool(raw.pending_nexus_operations)), + ) + if present + ] + assert not pending, ( + f"Workflow {handle.id} still has pending {', '.join(pending)} after {timeout}" + ) + + await assert_eventually(check, timeout=timeout, interval=interval) + + async def assert_task_fail_eventually( handle: WorkflowHandle, *, message_contains: str | None = None ) -> None: