49 lines
1.4 KiB
Python
49 lines
1.4 KiB
Python
|
"""Locust 性能压测脚本 — 天工智能体平台
|
||
|
|
启动: locust -f deploy/locustfile.py --host=http://localhost:8037
|
||
|
|
"""
|
||
|
|
from locust import HttpUser, task, between
|
||
|
|
import json
|
||
|
|
|
||
|
|
|
||
|
|
class AIPlatformUser(HttpUser):
|
||
|
|
wait_time = between(1, 3)
|
||
|
|
|
||
|
|
def on_start(self):
|
||
|
|
"""登录获取 token。"""
|
||
|
|
resp = self.client.post(
|
||
|
|
"/api/v1/auth/login",
|
||
|
|
headers={"Content-Type": "application/x-www-form-urlencoded"},
|
||
|
|
data={"username": "admin", "password": "123456"},
|
||
|
|
)
|
||
|
|
if resp.status_code == 200:
|
||
|
|
self.token = resp.json().get("access_token", "")
|
||
|
|
else:
|
||
|
|
self.token = ""
|
||
|
|
|
||
|
|
def _headers(self):
|
||
|
|
return {"Authorization": f"Bearer {self.token}"} if self.token else {}
|
||
|
|
|
||
|
|
@task(5)
|
||
|
|
def get_agents(self):
|
||
|
|
self.client.get("/api/v1/agents", headers=self._headers())
|
||
|
|
|
||
|
|
@task(3)
|
||
|
|
def agent_chat(self):
|
||
|
|
self.client.post(
|
||
|
|
"/api/v1/agent-chat/bare",
|
||
|
|
json={"message": "你好，请简单介绍一下自己"},
|
||
|
|
headers=self._headers(),
|
||
|
|
)
|
||
|
|
|
||
|
|
@task(2)
|
||
|
|
def get_executions(self):
|
||
|
|
self.client.get("/api/v1/executions?limit=10", headers=self._headers())
|
||
|
|
|
||
|
|
@task(2)
|
||
|
|
def health_check(self):
|
||
|
|
self.client.get("/health")
|
||
|
|
|
||
|
|
@task(1)
|
||
|
|
def get_workflows(self):
|
||
|
|
self.client.get("/api/v1/workflows?limit=20", headers=self._headers())
|