146 lines
5.4 KiB
Python
146 lines
5.4 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import statistics
|
|
import threading
|
|
import time
|
|
import urllib.error
|
|
import urllib.request
|
|
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description="Run a lightweight chat SSE load test.")
|
|
parser.add_argument("--base-url", default="http://127.0.0.1:8100/api")
|
|
parser.add_argument("--phone", required=True)
|
|
parser.add_argument("--code", default="123456")
|
|
parser.add_argument("--message", default="请用一句话回复：压测验证")
|
|
parser.add_argument("--concurrency", type=int, default=5)
|
|
parser.add_argument("--requests", type=int, default=20)
|
|
parser.add_argument("--timeout", type=int, default=90)
|
|
args = parser.parse_args()
|
|
|
|
token = login(args.base_url, args.phone, args.code, args.timeout)
|
|
lock = threading.Lock()
|
|
results: list[dict] = []
|
|
|
|
started_at = time.perf_counter()
|
|
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
|
|
futures = [
|
|
executor.submit(run_one, args.base_url, token, args.message, args.timeout)
|
|
for _ in range(args.requests)
|
|
]
|
|
for future in as_completed(futures):
|
|
result = future.result()
|
|
with lock:
|
|
results.append(result)
|
|
print(json.dumps(result, ensure_ascii=False))
|
|
|
|
elapsed = time.perf_counter() - started_at
|
|
print_summary(results, elapsed)
|
|
|
|
|
|
def login(base_url: str, phone: str, code: str, timeout: int) -> str:
|
|
body = post_json(f"{base_url}/auth/login", {"phone": phone, "code": code}, timeout=timeout)
|
|
return body["data"]["token"]
|
|
|
|
|
|
def run_one(base_url: str, token: str, message: str, timeout: int) -> dict:
|
|
started_at = time.perf_counter()
|
|
try:
|
|
session_body = post_json(
|
|
f"{base_url}/chat/session",
|
|
{"title": "压测会话"},
|
|
token=token,
|
|
timeout=timeout,
|
|
)
|
|
session_id = session_body["data"]["sessionId"]
|
|
first_event_ms = None
|
|
content_events = 0
|
|
error_message = ""
|
|
|
|
request = urllib.request.Request(
|
|
f"{base_url}/chat/completions",
|
|
data=json.dumps({"sessionId": session_id, "message": message}).encode("utf-8"),
|
|
headers={
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {token}",
|
|
},
|
|
)
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
for raw_line in response:
|
|
line = raw_line.decode("utf-8").strip()
|
|
if not line.startswith("data:"):
|
|
continue
|
|
if first_event_ms is None:
|
|
first_event_ms = elapsed_ms(started_at)
|
|
if line == "data: [DONE]":
|
|
break
|
|
payload = json.loads(line[5:].strip())
|
|
if payload.get("type") == "content":
|
|
content_events += 1
|
|
if payload.get("type") == "error":
|
|
error_message = payload.get("message", "unknown error")
|
|
|
|
return {
|
|
"ok": not error_message and content_events > 0,
|
|
"sessionId": session_id,
|
|
"firstEventMs": first_event_ms,
|
|
"totalMs": elapsed_ms(started_at),
|
|
"contentEvents": content_events,
|
|
"error": error_message,
|
|
}
|
|
except (urllib.error.URLError, TimeoutError, KeyError, ValueError) as exc:
|
|
return {
|
|
"ok": False,
|
|
"sessionId": None,
|
|
"firstEventMs": None,
|
|
"totalMs": elapsed_ms(started_at),
|
|
"contentEvents": 0,
|
|
"error": str(exc),
|
|
}
|
|
|
|
|
|
def post_json(url: str, payload: dict, *, token: str | None = None, timeout: int) -> dict:
|
|
headers = {"Content-Type": "application/json"}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
request = urllib.request.Request(url, data=json.dumps(payload).encode("utf-8"), headers=headers)
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
return json.loads(response.read().decode("utf-8"))
|
|
|
|
|
|
def print_summary(results: list[dict], elapsed: float) -> None:
|
|
ok_results = [result for result in results if result["ok"]]
|
|
failed_results = [result for result in results if not result["ok"]]
|
|
first_event_values = [result["firstEventMs"] for result in ok_results if result["firstEventMs"] is not None]
|
|
total_values = [result["totalMs"] for result in ok_results]
|
|
summary = {
|
|
"totalRequests": len(results),
|
|
"success": len(ok_results),
|
|
"failed": len(failed_results),
|
|
"elapsedSeconds": round(elapsed, 2),
|
|
"firstEventMsAvg": round(statistics.mean(first_event_values), 2) if first_event_values else None,
|
|
"firstEventMsP95": percentile(first_event_values, 95),
|
|
"totalMsAvg": round(statistics.mean(total_values), 2) if total_values else None,
|
|
"totalMsP95": percentile(total_values, 95),
|
|
}
|
|
print("SUMMARY " + json.dumps(summary, ensure_ascii=False))
|
|
|
|
|
|
def percentile(values: list[int], percentile_value: int) -> int | None:
|
|
if not values:
|
|
return None
|
|
sorted_values = sorted(values)
|
|
index = min(len(sorted_values) - 1, round((percentile_value / 100) * (len(sorted_values) - 1)))
|
|
return sorted_values[index]
|
|
|
|
|
|
def elapsed_ms(started_at: float) -> int:
|
|
return int((time.perf_counter() - started_at) * 1000)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|