From 9aac3399e76ff686ad864a76282784fdc1532992 Mon Sep 17 00:00:00 2001
From: Pavel Zorin <pazonec@yandex.ru>
Date: Wed, 19 Nov 2025 18:22:39 +0100
Subject: [PATCH] Fix Server Startup test

---
 .gitignore                               |  2 ++
 cognee/tests/test_cognee_server_start.py | 26 ++++++++++++++++--------
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/.gitignore b/.gitignore
index ff42efc97..1a330b040 100644
--- a/.gitignore
+++ b/.gitignore
@@ -13,6 +13,8 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+mise.toml
+
 full_run.ipynb
 
 # C extensions
diff --git a/cognee/tests/test_cognee_server_start.py b/cognee/tests/test_cognee_server_start.py
index ab68a8ef1..1bc819485 100644
--- a/cognee/tests/test_cognee_server_start.py
+++ b/cognee/tests/test_cognee_server_start.py
@@ -1,3 +1,4 @@
+from logging import getLogger
 import unittest
 import subprocess
 import time
@@ -8,6 +9,10 @@ from pathlib import Path
 import sys
 import uuid
 
+from tenacity import retry, stop_after_attempt, wait_fixed
+
+from cognee.api.health import HealthStatus
+
 
 class TestCogneeServerStart(unittest.TestCase):
     @classmethod
@@ -27,9 +32,8 @@ class TestCogneeServerStart(unittest.TestCase):
             stdout=subprocess.PIPE,
             stderr=subprocess.PIPE,
             preexec_fn=os.setsid,
-        )
-        # Give the server some time to start
-        time.sleep(35)
+        )        
+        cls.wait_for_server_startup()
 
         # Check if server started with errors
         if cls.server_process.poll() is not None:
@@ -49,12 +53,18 @@ class TestCogneeServerStart(unittest.TestCase):
                 cls.server_process.terminate()
             cls.server_process.wait()
 
+    @classmethod
+    @retry(wait=wait_fixed(2), stop=stop_after_attempt(5))
+    def wait_for_server_startup(cls):        
+        response = requests.get("http://localhost:8000/health", timeout=15)
+        status_str = response.json().get("status")
+        if response.status_code == 200 and status_str == "ready":
+            return True
+        else:
+            getLogger().info(f"Health check status {response.status_code} : {status_str}, retrying...")
+            raise requests.exceptions.RequestException(f"Server is {status_str}")
+        
     def test_server_is_running(self):
-        """Test that the server is running and can accept connections."""
-        # Test health endpoint
-        health_response = requests.get("http://localhost:8000/health", timeout=15)
-        self.assertIn(health_response.status_code, [200])
-
         # Test root endpoint
         root_response = requests.get("http://localhost:8000/", timeout=15)
         self.assertEqual(root_response.status_code, 200)