Merge pull request #697 from danielaskdd/improve-unit-test

Add http status check in unit tests.

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)}
@@ -742,7 +748,7 @@ Configuration file (config.json):
         nargs="+",
         choices=list(get_test_cases().keys()) + ["all"],
         default=["all"],
-        help="Test cases to run, options: %(choices)s. Use 'all' to run all tests",
+        help="Test cases to run, options: %(choices)s. Use 'all' to run all tests （except error tests)",
     )
     return parser.parse_args()
 
@@ -766,21 +772,18 @@ if __name__ == "__main__":
 
     try:
         if "all" in args.tests:
-            # Run all tests
+            # Run all tests except error handling tests
             if OutputControl.is_verbose():
                 print("\n【Chat API Tests】")
             run_test(test_non_stream_chat, "Non-streaming Chat Test")
             run_test(test_stream_chat, "Streaming Chat Test")
             run_test(test_query_modes, "Chat Query Mode Test")
-            run_test(test_error_handling, "Chat Error Handling Test")
-            run_test(test_stream_error_handling, "Chat Streaming Error Test")
 
             if OutputControl.is_verbose():
                 print("\n【Generate API Tests】")
             run_test(test_non_stream_generate, "Non-streaming Generate Test")
             run_test(test_stream_generate, "Streaming Generate Test")
             run_test(test_generate_with_system, "Generate with System Prompt Test")
-            run_test(test_generate_error_handling, "Generate Error Handling Test")
             run_test(test_generate_concurrent, "Generate Concurrent Test")
         else:
             # Run specified tests