fix(workflow): Disable sleep in message sample during pytest

boyangsvl · boyangsvl · commit 1936d464e12f · 2026-04-05T00:03:03.000-07:00
Change-Id: I561cd00fcbc132bb3198f6ffb02b7938f5eeda25
diff --git a/contributing/workflow_samples/message/agent.py b/contributing/workflow_samples/message/agent.py
@@ -14,13 +14,19 @@
 
 import asyncio
 import base64
+import os
 from typing import Any
 
 from google.adk import Event
 from google.adk import Workflow
 from google.genai import types
 
 
+async def sleep_if_not_pytest(seconds: float):
+  if "PYTEST_CURRENT_TEST" not in os.environ:
+    await asyncio.sleep(seconds)
+
+
 def send_string(node_input: Any = None):
   """Sends a single string message."""
   yield Event(message="#1 This is a simple string message.")
@@ -48,13 +54,13 @@ def send_multimodal(node_input: Any = None):
 async def multiple_messages(node_input: Any = None):
   """Sends multiple complete messages from the same node with an interval."""
   yield Event(message="#3 Multiple messages")
-  await asyncio.sleep(1.0)
+  await sleep_if_not_pytest(1.0)
 
   yield Event(message="Processing step 1...")
-  await asyncio.sleep(1.0)
+  await sleep_if_not_pytest(1.0)
 
   yield Event(message="Processing step 2...")
-  await asyncio.sleep(1.0)
+  await sleep_if_not_pytest(1.0)
 
   yield Event(message="Done processing.")
 
@@ -79,7 +85,7 @@ async def stream_sentence(node_input: Any = None):
   for i in range(0, len(sentence), 5):
     chunk = sentence[i : i + 5]
     yield Event(message=chunk, partial=True)
-    await asyncio.sleep(0.2)
+    await sleep_if_not_pytest(0.2)
 
 
 root_agent = Workflow(
