Skip to content

Commit 1d93565

Browse files
authored
[CE] Add base test class for web server testing (#3120)
* add test base class * fix codestyle * fix codestyle
1 parent e1011e9 commit 1d93565

5 files changed

Lines changed: 263 additions & 0 deletions

File tree

test/ce/server/core/__init__.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
#!/bin/env python3
2+
# -*- coding: utf-8 -*-
3+
# @author DDDivano
4+
# encoding=utf-8 vi:ts=4:sw=4:expandtab:ft=python
5+
import os
6+
import sys
7+
8+
from .logger import Logger
9+
10+
base_logger = Logger(loggername="FDSentry", save_level="channel", log_path="./fd_logs").get_logger()
11+
base_logger.setLevel("INFO")
12+
from .request_template import TEMPLATES
13+
from .utils import build_request_payload, send_request
14+
15+
__all__ = ["build_request_payload", "send_request", "TEMPLATES"]
16+
17+
# 检查环境变量是否存在
18+
URL = os.environ.get("URL")
19+
TEMPLATE = os.environ.get("TEMPLATE")
20+
21+
missing_vars = []
22+
if not URL:
23+
missing_vars.append("URL")
24+
if not TEMPLATE:
25+
missing_vars.append("TEMPLATE")
26+
27+
if missing_vars:
28+
msg = (
29+
f"❌ 缺少环境变量:{', '.join(missing_vars)},请先设置,例如:\n"
30+
f" export URL=http://localhost:8000/v1/chat/completions\n"
31+
f" export TEMPLATE=TOKEN_LOGPROB"
32+
)
33+
base_logger.error(msg)
34+
sys.exit(1) # 终止程序

test/ce/server/core/logger.py

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
#!/bin/env python3
2+
# -*- coding: utf-8 -*-
3+
# @author DDDivano
4+
# encoding=utf-8 vi:ts=4:sw=4:expandtab:ft=python
5+
"""
6+
ServeTest
7+
"""
8+
import logging
9+
import os
10+
from datetime import datetime
11+
12+
import pytz
13+
14+
15+
class Logger(object):
16+
"""
17+
日志记录配置的基础类。
18+
"""
19+
20+
SAVE_LEVELS = ["both", "file", "channel"]
21+
LOG_FORMAT = "%(asctime)s - %(name)s - [%(levelname)s] - %(message)s"
22+
23+
def __init__(self, loggername, save_level="both", log_path=None):
24+
"""
25+
使用指定名称和保存级别初始化日志记录器。
26+
27+
Args:
28+
loggername (str): 日志记录器的名称。
29+
save_level (str): 日志保存的级别。默认为"both"。file: 仅保存到文件,channel: 仅保存到控制台。
30+
log_path (str, optional): 日志文件保存路径。默认为None。
31+
"""
32+
33+
if save_level not in self.SAVE_LEVELS:
34+
raise ValueError(f"Invalid save level: {save_level}. Allowed values: {self.SAVE_LEVELS}")
35+
36+
self.logger = logging.getLogger(loggername)
37+
self.logger.setLevel(logging.DEBUG)
38+
39+
# 设置时区为东八区
40+
tz = pytz.timezone("Asia/Shanghai")
41+
42+
# 自定义时间格式化器,指定时区为东八区
43+
class CSTFormatter(logging.Formatter):
44+
"""
45+
自定义时间格式化器,指定时区为东八区
46+
"""
47+
48+
def converter(self, timestamp):
49+
"""
50+
自定义时间转换函数,加上时区信息
51+
Args:
52+
timestamp (int): 时间戳。
53+
Returns:
54+
tuple: 格式化后的时间元组。
55+
"""
56+
dt = datetime.utcfromtimestamp(timestamp)
57+
dt = pytz.utc.localize(dt).astimezone(tz)
58+
return dt.timetuple()
59+
60+
formatter = CSTFormatter(self.LOG_FORMAT)
61+
log_name = None
62+
if save_level == "both" or save_level == "file":
63+
os.makedirs(log_path, exist_ok=True)
64+
log_filename = f"out_{datetime.now().strftime('%Y-%m-%d_%H-%M-%S')}.log"
65+
log_name = os.path.join(log_path, log_filename)
66+
file_handler = logging.FileHandler(log_name, encoding="utf-8")
67+
file_handler.setLevel(logging.DEBUG)
68+
file_handler.setFormatter(formatter)
69+
self.logger.addHandler(file_handler)
70+
71+
if save_level == "both" or save_level == "channel":
72+
console_handler = logging.StreamHandler()
73+
console_handler.setLevel(logging.DEBUG)
74+
console_handler.setFormatter(formatter)
75+
self.logger.addHandler(console_handler)
76+
77+
if log_name is None:
78+
self.logger.info(
79+
f"Logger initialized. Log level: {save_level}. "
80+
f"Log path ({log_path}) is unused according to the level."
81+
)
82+
else:
83+
self.logger.info(f"Logger initialized. Log level: {save_level}. Log path: {log_name}")
84+
# Adjusting the timezone offset
85+
86+
def get_logger(self):
87+
"""
88+
Get the logger object
89+
"""
90+
return self.logger
91+
92+
93+
if __name__ == "__main__":
94+
# Test the logger
95+
logger = Logger("test_logger", save_level="channel").get_logger()
96+
logger.info("the is the beginning")
97+
logger.debug("the is the beginning")
98+
logger.warning("the is the beginning")
99+
logger.error("the is the beginning")
Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
#!/bin/env python3
2+
# -*- coding: utf-8 -*-
3+
# @author DDDivano
4+
# encoding=utf-8 vi:ts=4:sw=4:expandtab:ft=python
5+
"""
6+
ServeTest
7+
"""
8+
9+
10+
TOKEN_LOGPROB = {
11+
"model": "default",
12+
"temperature": 0,
13+
"top_p": 0,
14+
"seed": 33,
15+
"stream": True,
16+
"logprobs": True,
17+
"top_logprobs": 5,
18+
"max_tokens": 10000,
19+
}
20+
21+
22+
TEMPLATES = {
23+
"TOKEN_LOGPROB": TOKEN_LOGPROB,
24+
# "ANOTHER_TEMPLATE": ANOTHER_TEMPLATE
25+
}

test/ce/server/core/utils.py

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
#!/bin/env python3
2+
# -*- coding: utf-8 -*-
3+
# @author DDDivano
4+
# encoding=utf-8 vi:ts=4:sw=4:expandtab:ft=python
5+
6+
import requests
7+
from core import TEMPLATES, base_logger
8+
9+
10+
def build_request_payload(template_name: str, case_data: dict) -> dict:
11+
"""
12+
基于模板构造请求 payload,按优先级依次合并:
13+
template < payload 参数 < case_data,后者会覆盖前者的同名字段。
14+
15+
:param template_name: 模板变量名,例如 "TOKEN_LOGPROB"
16+
:return: 构造后的完整请求 payload dict
17+
"""
18+
template = TEMPLATES[template_name]
19+
print(template)
20+
final_payload = template.copy()
21+
final_payload.update(case_data)
22+
23+
return final_payload
24+
25+
26+
def send_request(url, payload, timeout=600, stream=False):
27+
"""
28+
向指定URL发送POST请求,并返回响应结果。
29+
30+
Args:
31+
url (str): 请求的目标URL。
32+
payload (dict): 请求的负载数据,应该是一个字典类型。
33+
timeout (int, optional): 请求的超时时间,默认为600秒。
34+
stream (bool, optional): 是否以流的方式下载响应内容,默认为False。
35+
36+
Returns:
37+
response: 请求的响应结果,如果请求失败则返回None。
38+
39+
Raises:
40+
None
41+
42+
"""
43+
headers = {
44+
"Content-Type": "application/json",
45+
}
46+
base_logger.info("🔄 正在请求模型接口...")
47+
48+
try:
49+
res = requests.post(url, headers=headers, json=payload, stream=stream, timeout=timeout)
50+
base_logger.info("🟢 接收响应中...\n")
51+
return res
52+
except requests.exceptions.Timeout:
53+
base_logger.error(f"❌ 请求超时(超过 {timeout} 秒)")
54+
return None
55+
except requests.exceptions.RequestException as e:
56+
base_logger.error(f"❌ 请求失败:{e}")
57+
return None

test/ce/server/demo.py

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
1+
#!/bin/env python3
2+
# -*- coding: utf-8 -*-
3+
# @author DDDivano
4+
# encoding=utf-8 vi:ts=4:sw=4:expandtab:ft=python
5+
6+
from core import TEMPLATE, URL, build_request_payload, send_request
7+
8+
9+
def demo():
10+
data = {
11+
"stream": False,
12+
"messages": [
13+
{"role": "system", "content": "You are a helpful assistant."},
14+
{"role": "user", "content": "牛顿的三大运动定律是什么?"},
15+
],
16+
"max_tokens": 3,
17+
}
18+
payload = build_request_payload(TEMPLATE, data)
19+
req = send_request(URL, payload)
20+
print(req.json())
21+
req = req.json()
22+
23+
assert req["usage"]["prompt_tokens"] == 22
24+
assert req["usage"]["total_tokens"] == 25
25+
assert req["usage"]["completion_tokens"] == 3
26+
27+
28+
def test_demo():
29+
data = {
30+
"stream": False,
31+
"messages": [
32+
{"role": "system", "content": "You are a helpful assistant."},
33+
{"role": "user", "content": "牛顿的三大运动定律是什么?"},
34+
],
35+
"max_tokens": 3,
36+
}
37+
payload = build_request_payload(TEMPLATE, data)
38+
req = send_request(URL, payload)
39+
print(req.json())
40+
req = req.json()
41+
42+
assert req["usage"]["prompt_tokens"] == 22
43+
assert req["usage"]["total_tokens"] == 25
44+
assert req["usage"]["completion_tokens"] == 5
45+
46+
47+
if __name__ == "__main__":
48+
demo()

0 commit comments

Comments
 (0)