From c316edb2d990c66e9e73aabba5c46c3c21902181 Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Sun, 2 Feb 2025 22:03:55 +0800
Subject: [PATCH] Add http status check for unit tests

---
 test_lightrag_ollama_chat.py | 24 +++++++++++++++---------
 1 file changed, 15 insertions(+), 9 deletions(-)

diff --git a/test_lightrag_ollama_chat.py b/test_lightrag_ollama_chat.py
index f4ee42c0..d1375c8c 100644
--- a/test_lightrag_ollama_chat.py
+++ b/test_lightrag_ollama_chat.py
@@ -123,18 +123,20 @@ class TestStats:
 
 
 def make_request(
-    url: str, data: Dict[str, Any], stream: bool = False
+    url: str, data: Dict[str, Any], stream: bool = False, check_status: bool = True
 ) -> requests.Response:
     """Send an HTTP request with retry mechanism
     Args:
         url: Request URL
         data: Request data
         stream: Whether to use streaming response
+        check_status: Whether to check HTTP status code (default: True)
     Returns:
         requests.Response: Response object
 
     Raises:
         requests.exceptions.RequestException: Request failed after all retries
+        requests.exceptions.HTTPError: HTTP status code is not 200 (when check_status is True)
     """
     server_config = CONFIG["server"]
     max_retries = server_config["max_retries"]
@@ -144,6 +146,8 @@ def make_request(
     for attempt in range(max_retries):
         try:
             response = requests.post(url, json=data, stream=stream, timeout=timeout)
+            if check_status and response.status_code != 200:
+                response.raise_for_status()
             return response
         except requests.exceptions.RequestException as e:
             if attempt == max_retries - 1:  # Last retry
@@ -433,7 +437,7 @@ def test_stream_error_handling() -> None:
     if OutputControl.is_verbose():
         print("\n--- Testing empty message list (streaming) ---")
     data = create_error_test_data("empty_messages")
-    response = make_request(url, data, stream=True)
+    response = make_request(url, data, stream=True, check_status=False)
     print(f"Status code: {response.status_code}")
     if response.status_code != 200:
         print_json_response(response.json(), "Error message")
@@ -443,7 +447,7 @@ def test_stream_error_handling() -> None:
     if OutputControl.is_verbose():
         print("\n--- Testing invalid role field (streaming) ---")
     data = create_error_test_data("invalid_role")
-    response = make_request(url, data, stream=True)
+    response = make_request(url, data, stream=True, check_status=False)
     print(f"Status code: {response.status_code}")
     if response.status_code != 200:
         print_json_response(response.json(), "Error message")
@@ -453,7 +457,7 @@ def test_stream_error_handling() -> None:
     if OutputControl.is_verbose():
         print("\n--- Testing missing content field (streaming) ---")
     data = create_error_test_data("missing_content")
-    response = make_request(url, data, stream=True)
+    response = make_request(url, data, stream=True, check_status=False)
     print(f"Status code: {response.status_code}")
     if response.status_code != 200:
         print_json_response(response.json(), "Error message")
@@ -484,7 +488,7 @@ def test_error_handling() -> None:
         print("\n--- Testing empty message list ---")
     data = create_error_test_data("empty_messages")
     data["stream"] = False  # Change to non-streaming mode
-    response = make_request(url, data)
+    response = make_request(url, data, check_status=False)
     print(f"Status code: {response.status_code}")
     print_json_response(response.json(), "Error message")
 
@@ -493,7 +497,7 @@ def test_error_handling() -> None:
         print("\n--- Testing invalid role field ---")
     data = create_error_test_data("invalid_role")
     data["stream"] = False  # Change to non-streaming mode
-    response = make_request(url, data)
+    response = make_request(url, data, check_status=False)
     print(f"Status code: {response.status_code}")
     print_json_response(response.json(), "Error message")
 
@@ -502,7 +506,7 @@ def test_error_handling() -> None:
         print("\n--- Testing missing content field ---")
     data = create_error_test_data("missing_content")
     data["stream"] = False  # Change to non-streaming mode
-    response = make_request(url, data)
+    response = make_request(url, data, check_status=False)
     print(f"Status code: {response.status_code}")
     print_json_response(response.json(), "Error message")
 
@@ -609,7 +613,7 @@ def test_generate_error_handling() -> None:
     if OutputControl.is_verbose():
         print("\n=== Testing empty prompt ===")
     data = create_generate_request_data("", stream=False)
-    response = make_request(url, data)
+    response = make_request(url, data, check_status=False)
     print(f"Status code: {response.status_code}")
     print_json_response(response.json(), "Error message")
 
@@ -621,7 +625,7 @@ def test_generate_error_handling() -> None:
         options={"invalid_option": "value"},
         stream=False,
     )
-    response = make_request(url, data)
+    response = make_request(url, data, check_status=False)
     print(f"Status code: {response.status_code}")
     print_json_response(response.json(), "Error message")
 
@@ -642,6 +646,8 @@ def test_generate_concurrent() -> None:
         data = create_generate_request_data(prompt, stream=False)
         try:
             async with session.post(url, json=data) as response:
+                if response.status != 200:
+                    response.raise_for_status()
                 return await response.json()
         except Exception as e:
             return {"error": str(e)}