yangdx commited on
Commit
30c0560
·
1 Parent(s): 153eff6

Add http status check for unit tests

Browse files
Files changed (1) hide show
  1. test_lightrag_ollama_chat.py +15 -9
test_lightrag_ollama_chat.py CHANGED
@@ -123,18 +123,20 @@ class TestStats:
123
 
124
 
125
  def make_request(
126
- url: str, data: Dict[str, Any], stream: bool = False
127
  ) -> requests.Response:
128
  """Send an HTTP request with retry mechanism
129
  Args:
130
  url: Request URL
131
  data: Request data
132
  stream: Whether to use streaming response
 
133
  Returns:
134
  requests.Response: Response object
135
 
136
  Raises:
137
  requests.exceptions.RequestException: Request failed after all retries
 
138
  """
139
  server_config = CONFIG["server"]
140
  max_retries = server_config["max_retries"]
@@ -144,6 +146,8 @@ def make_request(
144
  for attempt in range(max_retries):
145
  try:
146
  response = requests.post(url, json=data, stream=stream, timeout=timeout)
 
 
147
  return response
148
  except requests.exceptions.RequestException as e:
149
  if attempt == max_retries - 1: # Last retry
@@ -433,7 +437,7 @@ def test_stream_error_handling() -> None:
433
  if OutputControl.is_verbose():
434
  print("\n--- Testing empty message list (streaming) ---")
435
  data = create_error_test_data("empty_messages")
436
- response = make_request(url, data, stream=True)
437
  print(f"Status code: {response.status_code}")
438
  if response.status_code != 200:
439
  print_json_response(response.json(), "Error message")
@@ -443,7 +447,7 @@ def test_stream_error_handling() -> None:
443
  if OutputControl.is_verbose():
444
  print("\n--- Testing invalid role field (streaming) ---")
445
  data = create_error_test_data("invalid_role")
446
- response = make_request(url, data, stream=True)
447
  print(f"Status code: {response.status_code}")
448
  if response.status_code != 200:
449
  print_json_response(response.json(), "Error message")
@@ -453,7 +457,7 @@ def test_stream_error_handling() -> None:
453
  if OutputControl.is_verbose():
454
  print("\n--- Testing missing content field (streaming) ---")
455
  data = create_error_test_data("missing_content")
456
- response = make_request(url, data, stream=True)
457
  print(f"Status code: {response.status_code}")
458
  if response.status_code != 200:
459
  print_json_response(response.json(), "Error message")
@@ -484,7 +488,7 @@ def test_error_handling() -> None:
484
  print("\n--- Testing empty message list ---")
485
  data = create_error_test_data("empty_messages")
486
  data["stream"] = False # Change to non-streaming mode
487
- response = make_request(url, data)
488
  print(f"Status code: {response.status_code}")
489
  print_json_response(response.json(), "Error message")
490
 
@@ -493,7 +497,7 @@ def test_error_handling() -> None:
493
  print("\n--- Testing invalid role field ---")
494
  data = create_error_test_data("invalid_role")
495
  data["stream"] = False # Change to non-streaming mode
496
- response = make_request(url, data)
497
  print(f"Status code: {response.status_code}")
498
  print_json_response(response.json(), "Error message")
499
 
@@ -502,7 +506,7 @@ def test_error_handling() -> None:
502
  print("\n--- Testing missing content field ---")
503
  data = create_error_test_data("missing_content")
504
  data["stream"] = False # Change to non-streaming mode
505
- response = make_request(url, data)
506
  print(f"Status code: {response.status_code}")
507
  print_json_response(response.json(), "Error message")
508
 
@@ -609,7 +613,7 @@ def test_generate_error_handling() -> None:
609
  if OutputControl.is_verbose():
610
  print("\n=== Testing empty prompt ===")
611
  data = create_generate_request_data("", stream=False)
612
- response = make_request(url, data)
613
  print(f"Status code: {response.status_code}")
614
  print_json_response(response.json(), "Error message")
615
 
@@ -621,7 +625,7 @@ def test_generate_error_handling() -> None:
621
  options={"invalid_option": "value"},
622
  stream=False,
623
  )
624
- response = make_request(url, data)
625
  print(f"Status code: {response.status_code}")
626
  print_json_response(response.json(), "Error message")
627
 
@@ -642,6 +646,8 @@ def test_generate_concurrent() -> None:
642
  data = create_generate_request_data(prompt, stream=False)
643
  try:
644
  async with session.post(url, json=data) as response:
 
 
645
  return await response.json()
646
  except Exception as e:
647
  return {"error": str(e)}
 
123
 
124
 
125
  def make_request(
126
+ url: str, data: Dict[str, Any], stream: bool = False, check_status: bool = True
127
  ) -> requests.Response:
128
  """Send an HTTP request with retry mechanism
129
  Args:
130
  url: Request URL
131
  data: Request data
132
  stream: Whether to use streaming response
133
+ check_status: Whether to check HTTP status code (default: True)
134
  Returns:
135
  requests.Response: Response object
136
 
137
  Raises:
138
  requests.exceptions.RequestException: Request failed after all retries
139
+ requests.exceptions.HTTPError: HTTP status code is not 200 (when check_status is True)
140
  """
141
  server_config = CONFIG["server"]
142
  max_retries = server_config["max_retries"]
 
146
  for attempt in range(max_retries):
147
  try:
148
  response = requests.post(url, json=data, stream=stream, timeout=timeout)
149
+ if check_status and response.status_code != 200:
150
+ response.raise_for_status()
151
  return response
152
  except requests.exceptions.RequestException as e:
153
  if attempt == max_retries - 1: # Last retry
 
437
  if OutputControl.is_verbose():
438
  print("\n--- Testing empty message list (streaming) ---")
439
  data = create_error_test_data("empty_messages")
440
+ response = make_request(url, data, stream=True, check_status=False)
441
  print(f"Status code: {response.status_code}")
442
  if response.status_code != 200:
443
  print_json_response(response.json(), "Error message")
 
447
  if OutputControl.is_verbose():
448
  print("\n--- Testing invalid role field (streaming) ---")
449
  data = create_error_test_data("invalid_role")
450
+ response = make_request(url, data, stream=True, check_status=False)
451
  print(f"Status code: {response.status_code}")
452
  if response.status_code != 200:
453
  print_json_response(response.json(), "Error message")
 
457
  if OutputControl.is_verbose():
458
  print("\n--- Testing missing content field (streaming) ---")
459
  data = create_error_test_data("missing_content")
460
+ response = make_request(url, data, stream=True, check_status=False)
461
  print(f"Status code: {response.status_code}")
462
  if response.status_code != 200:
463
  print_json_response(response.json(), "Error message")
 
488
  print("\n--- Testing empty message list ---")
489
  data = create_error_test_data("empty_messages")
490
  data["stream"] = False # Change to non-streaming mode
491
+ response = make_request(url, data, check_status=False)
492
  print(f"Status code: {response.status_code}")
493
  print_json_response(response.json(), "Error message")
494
 
 
497
  print("\n--- Testing invalid role field ---")
498
  data = create_error_test_data("invalid_role")
499
  data["stream"] = False # Change to non-streaming mode
500
+ response = make_request(url, data, check_status=False)
501
  print(f"Status code: {response.status_code}")
502
  print_json_response(response.json(), "Error message")
503
 
 
506
  print("\n--- Testing missing content field ---")
507
  data = create_error_test_data("missing_content")
508
  data["stream"] = False # Change to non-streaming mode
509
+ response = make_request(url, data, check_status=False)
510
  print(f"Status code: {response.status_code}")
511
  print_json_response(response.json(), "Error message")
512
 
 
613
  if OutputControl.is_verbose():
614
  print("\n=== Testing empty prompt ===")
615
  data = create_generate_request_data("", stream=False)
616
+ response = make_request(url, data, check_status=False)
617
  print(f"Status code: {response.status_code}")
618
  print_json_response(response.json(), "Error message")
619
 
 
625
  options={"invalid_option": "value"},
626
  stream=False,
627
  )
628
+ response = make_request(url, data, check_status=False)
629
  print(f"Status code: {response.status_code}")
630
  print_json_response(response.json(), "Error message")
631
 
 
646
  data = create_generate_request_data(prompt, stream=False)
647
  try:
648
  async with session.post(url, json=data) as response:
649
+ if response.status != 200:
650
+ response.raise_for_status()
651
  return await response.json()
652
  except Exception as e:
653
  return {"error": str(e)}