138 lines
4.7 KiB
Python
138 lines
4.7 KiB
Python
from logging import getLogger
|
|
import unittest
|
|
import subprocess
|
|
import time
|
|
import os
|
|
import signal
|
|
import requests
|
|
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
|
|
def setUpClass(cls):
|
|
# Start the Cognee server - just check if the server can start without errors
|
|
cls.server_process = subprocess.Popen(
|
|
[
|
|
sys.executable,
|
|
"-m",
|
|
"uvicorn",
|
|
"cognee.api.client:app",
|
|
"--host",
|
|
"0.0.0.0",
|
|
"--port",
|
|
"8000",
|
|
],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.PIPE,
|
|
preexec_fn=os.setsid,
|
|
)
|
|
cls.wait_for_server_startup()
|
|
|
|
# Check if server started with errors
|
|
if cls.server_process.poll() is not None:
|
|
stderr = cls.server_process.stderr.read().decode("utf-8")
|
|
print(f"Server failed to start: {stderr}", file=sys.stderr)
|
|
raise RuntimeError(f"Server failed to start: {stderr}")
|
|
|
|
@classmethod
|
|
def tearDownClass(cls):
|
|
# Terminate the server process
|
|
if hasattr(cls, "server_process") and cls.server_process:
|
|
if hasattr(os, "killpg"):
|
|
# Unix-like systems: Use process groups
|
|
os.killpg(os.getpgid(cls.server_process.pid), signal.SIGTERM)
|
|
else:
|
|
# Windows: Just terminate the main process
|
|
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 root endpoint
|
|
root_response = requests.get("http://localhost:8000/", timeout=15)
|
|
self.assertEqual(root_response.status_code, 200)
|
|
self.assertIn("message", root_response.json())
|
|
self.assertEqual(root_response.json()["message"], "Hello, World, I am alive!")
|
|
|
|
# Login request
|
|
url = "http://127.0.0.1:8000/api/v1/auth/login"
|
|
form_data = {
|
|
"username": "default_user@example.com",
|
|
"password": "default_password",
|
|
}
|
|
login_response = requests.post(url, data=form_data, timeout=15)
|
|
login_response.raise_for_status() # raises on HTTP 4xx/5xx
|
|
|
|
# Define Bearer token to use for authorization
|
|
auth_var = (
|
|
"Bearer " + login_response.json()["access_token"]
|
|
) # e.g. "Bearer eyJhbGciOiJIUzI1NiIsInR5cCI6..."
|
|
|
|
# Add request
|
|
url = "http://127.0.0.1:8000/api/v1/add"
|
|
file_path = Path(os.path.join(Path(__file__).parent, "test_data/example.png"))
|
|
headers = {"Authorization": auth_var}
|
|
|
|
dataset_name = f"test_{uuid.uuid4().hex[:8]}"
|
|
form_data = {"datasetName": dataset_name}
|
|
|
|
file = {
|
|
"data": (
|
|
file_path.name,
|
|
open(file_path, "rb"),
|
|
)
|
|
}
|
|
|
|
payload = {"datasets": [dataset_name]}
|
|
|
|
add_response = requests.post(url, headers=headers, data=form_data, files=file, timeout=50)
|
|
if add_response.status_code not in [200, 201]:
|
|
add_response.raise_for_status()
|
|
|
|
# Cognify request
|
|
url = "http://127.0.0.1:8000/api/v1/cognify"
|
|
headers = {
|
|
"Authorization": auth_var,
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
cognify_response = requests.post(url, headers=headers, json=payload, timeout=150)
|
|
if cognify_response.status_code not in [200, 201]:
|
|
cognify_response.raise_for_status()
|
|
|
|
# TODO: Add test to verify cognify pipeline is complete before testing search
|
|
|
|
# Search request
|
|
url = "http://127.0.0.1:8000/api/v1/search"
|
|
|
|
headers = {
|
|
"Authorization": auth_var,
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
payload = {"searchType": "GRAPH_COMPLETION", "query": "What's in the document?"}
|
|
|
|
search_response = requests.post(url, headers=headers, json=payload, timeout=50)
|
|
if search_response.status_code not in [200, 201]:
|
|
search_response.raise_for_status()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|