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")