Commit b1cde70d authored by root's avatar root

fix: use correct mock API paths in locust

parent f8dec13f
"""
Locust Load Test - Canifa Chatbot API
Run: locust -f locustfile.py --host=http://localhost:5000
"""
import random
from locust import between, events, task
from locust.contrib.fasthttp import FastHttpUser
# ============ CUSTOM UI PARAMETER ============
@events.init_command_line_parser.add_listener
def _(parser):
parser.add_argument(
"--api-target",
choices=["All", "MockAgentChat", "RealDBSearch", "RealRetrieverDB"],
default="All",
help="Chọn API cụ thể để stress test",
)
# ============ SAMPLE TEST DATA ============
SAMPLE_QUERIES = [
"Tìm áo sơ mi nam công sở",
"Quần jeans nữ skinny",
"Áo khoác da nam cao cấp",
"Đầm dự tiệc sang trọng",
"Giày thể thao chạy bộ",
"Túi xách nữ da thật",
"Áo thun nam basic",
"Váy công sở nữ",
"Áo len nữ mùa đông",
"Giày cao gót nữ 7cm",
"Quần short nam mùa hè",
"Balo laptop chống nước",
"Áo polo nam trơn",
"Chân váy bút chì",
"Áo hoodie unisex",
]
SAMPLE_MAGENTO_CODES = [
"10001001",
"10001002",
"10001003",
"10002001",
"10002002",
]
PRICE_RANGES = [
{"min": None, "max": None},
{"min": 100000, "max": 500000},
{"min": 500000, "max": 1000000},
{"min": 1000000, "max": 2000000},
]
# ============ MAIN LOAD TEST CLASS ============
class CanifaStressTestUser(FastHttpUser):
"""Single User Class with API Selection"""
# Người dùng đọc + gõ mất khoảng 15 đến 25 giây rồi mới gửi câu tiếp theo
wait_time = between(15, 25)
def on_start(self):
self.user_id = f"locust_user_{random.randint(1000, 9999)}"
self.session_id = f"session_{random.randint(10000, 99999)}"
# Lấy giá trị từ giao diện Web
self.api_target = self.environment.parsed_options.api_target
# --- LOGIC CÁC API ---
def run_mock_chat(self):
query = random.choice(SAMPLE_QUERIES)
payload = {
"user_query": query,
"user_id": self.user_id,
"session_id": self.session_id,
}
self.client.post(
"/api/mock/agent/chat", json=payload, name="[API] Mock Agent Chat"
)
def run_db_search(self):
query = random.choice(SAMPLE_QUERIES)
price_range = random.choice(PRICE_RANGES)
payload = {
"query": query,
"price_min": price_range["min"],
"price_max": price_range["max"],
"top_k": 10,
}
self.client.post(
"/api/mock/db/search", json=payload, name="[API] Real DB Search"
)
def run_retriever(self):
query = random.choice(SAMPLE_QUERIES)
payload = {
"user_query": query,
"user_id": self.user_id,
"session_id": self.session_id,
}
self.client.post(
"/api/mock/retrieverdb", json=payload, name="[API] Real Retriever"
)
# --- TASK DISPATCHER ---
@task
def dispatcher(self):
target = self.api_target
if target == "MockAgentChat":
self.run_mock_chat()
elif target == "RealDBSearch":
self.run_db_search()
elif target == "RealRetrieverDB":
self.run_retriever()
else: # "All"
# Random chọn 1 trong 3 theo tỷ lệ (3:2:1)
choice = random.choices(["chat", "search", "retriever"], weights=[3, 2, 1])[
0
]
if choice == "chat":
self.run_mock_chat()
elif choice == "search":
self.run_db_search()
else:
self.run_retriever()
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment