From 55c51f2d2840aa77a132b2518c3e6646a24088b0 Mon Sep 17 00:00:00 2001
From: Matthew Zhou <mattzh1314@gmail.com>
Date: Mon, 28 Apr 2025 15:23:53 -0700
Subject: [PATCH] test: Add wait for server startup for letta batch tests
 (#1916)

---
 tests/test_letta_agent_batch.py |  4 ++--
 tests/utils.py                  | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/tests/test_letta_agent_batch.py b/tests/test_letta_agent_batch.py
index 1cde5dc8..20f16611 100644
--- a/tests/test_letta_agent_batch.py
+++ b/tests/test_letta_agent_batch.py
@@ -1,6 +1,5 @@
 import os
 import threading
-import time
 from datetime import datetime, timezone
 from typing import Tuple
 from unittest.mock import AsyncMock, patch
@@ -32,6 +31,7 @@ from letta.schemas.message import MessageCreate
 from letta.schemas.tool_rule import InitToolRule
 from letta.server.db import db_context
 from letta.server.server import SyncServer
+from tests.utils import wait_for_server
 
 # --------------------------------------------------------------------------- #
 # Test Constants
@@ -311,7 +311,7 @@ def server_url():
     if not os.getenv("LETTA_SERVER_URL"):
         thread = threading.Thread(target=run_server, daemon=True)
         thread.start()
-        time.sleep(1)  # Give server time to start
+        wait_for_server(url)
 
     return url
 
diff --git a/tests/utils.py b/tests/utils.py
index 7ca72d5b..12f7c40a 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -175,3 +175,21 @@ def wait_for_incoming_message(
         time.sleep(sleep_interval)
 
     return False
+
+
+def wait_for_server(url, timeout=30, interval=0.5):
+    """Wait for server to become available by polling the given URL."""
+    import requests
+    from requests.exceptions import ConnectionError
+
+    start_time = time.time()
+    while time.time() - start_time < timeout:
+        try:
+            response = requests.get(f"{url}/v1/health", timeout=2)
+            if response.status_code == 200:
+                return True
+        except (ConnectionError, requests.Timeout):
+            pass
+        time.sleep(interval)
+
+    raise TimeoutError(f"Server at {url} did not start within {timeout} seconds")