Locust压测ai接口

62 阅读1分钟

安装miniconda然后启动虚拟环境 pip locust 测试chatflow代码 更改Authorization中的app-Jgx9XiL4q186TMHD8PJIqp7B成dify的api密钥

from locust import HttpUser, TaskSet, task, between
import time
 
class ChatMessages(TaskSet):
    @task
    def chat_messages(self):
        url = "/chat-messages"
        headers = {
            "Authorization": "Bearer app-Jgx9XiL4q186TMHD8PJIqp7B",
            "Content-Type": "application/json",
        }
        payload = {
                      "inputs": {},
                      "query": "压测",
                      "response_mode": "streaming",
                      "user": "压测"
                 }
 
  # 记录开始时间
        start_time = time.time()
    try:
            # 发起 POST 请求
            with self.client.post(url, json=payload, headers=headers, stream=True, catch_response=True, timeout=60) as response:
                if response.status_code != 200:
                    response.failure(f"Unexpected status code: {response.status_code}")
                    return
 
                # 记录首字节时间(TTFB)
                ttfb = time.time() - start_time
                print(f"TTFB: {ttfb:.2f}s")
                # 逐块读取响应内容
                chunk_count = 0
                for line in response.iter_lines(decode_unicode=True):
                    if line.startswith("data:"):
                        chunk_count += 1
                        total_time = time.time() - start_time
                        print(f"Received chunk #{chunk_count}: {line}{total_time:.2f}s")
                    if "message_end" in line:
                        total_time = time.time() - start_time
                        print(f"Total response time: {total_time:.2f}s")
                        break
        except Exception as e:
            response.failure(f"Request failed: {str(e)}")

class ChatMessagesTest(HttpUser):
    # 声明执行的任务集是哪个类
    tasks = [ChatMessages]
    # 设置运行过程中间隔时间
#     wait_time = between(1, 2)

    # 每用户结束动作,作用等同于pytest、unittest的teardown
    def on_stop(self):
        self.client.close()

启动 locust -f test.py 测试 ip:8089

image.png