140 lines
5.1 KiB
Python
140 lines
5.1 KiB
Python
import requests
|
||
import json
|
||
import logging
|
||
import time
|
||
import os
|
||
|
||
# 配置日志
|
||
logging.basicConfig(
|
||
level=logging.INFO,
|
||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||
)
|
||
logger = logging.getLogger("TestQWenImage")
|
||
|
||
# API基础URL
|
||
base_url = "http://localhost:8200/api/qwenImage"
|
||
|
||
|
||
def test_get_config():
|
||
"""测试获取配置接口"""
|
||
try:
|
||
url = f"{base_url}/config"
|
||
logger.info(f"调用获取配置接口: {url}")
|
||
response = requests.get(url)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
logger.info(f"获取配置成功: {json.dumps(result, ensure_ascii=False, indent=2)}")
|
||
return True, result
|
||
else:
|
||
logger.error(f"获取配置失败: HTTP状态码={response.status_code}, 响应内容={response.text}")
|
||
return False, None
|
||
except Exception as e:
|
||
logger.exception(f"获取配置时发生异常: {str(e)}")
|
||
return False, None
|
||
|
||
|
||
def test_generate_image(prompt, n=1, size='1328*1328', save_local=True):
|
||
"""测试生成图片接口
|
||
|
||
Args:
|
||
prompt: 图片描述提示词
|
||
n: 生成图片数量
|
||
size: 图片尺寸
|
||
save_local: 是否保存到本地
|
||
|
||
Returns:
|
||
tuple: (是否成功, 结果数据)
|
||
"""
|
||
try:
|
||
url = f"{base_url}/generate"
|
||
headers = {"Content-Type": "application/json"}
|
||
data = {
|
||
"prompt": prompt,
|
||
"n": n,
|
||
"size": size,
|
||
"save_local": save_local
|
||
}
|
||
|
||
logger.info(f"调用生成图片接口: {url}")
|
||
logger.info(f"请求参数: prompt={prompt[:50]}..., n={n}, size={size}, save_local={save_local}")
|
||
|
||
# 记录开始时间
|
||
start_time = time.time()
|
||
|
||
# 发送请求
|
||
response = requests.post(url, headers=headers, data=json.dumps(data, ensure_ascii=False))
|
||
|
||
# 计算耗时
|
||
elapsed_time = time.time() - start_time
|
||
logger.info(f"请求耗时: {elapsed_time:.2f}秒")
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
logger.info(f"生成图片成功: {json.dumps(result, ensure_ascii=False, indent=2)}")
|
||
|
||
# 检查返回的数据
|
||
if result.get("code") == 200 and "data" in result:
|
||
images = result["data"].get("images", [])
|
||
logger.info(f"成功生成{len(images)}张图片")
|
||
|
||
# 如果保存了本地文件，检查文件是否存在
|
||
if save_local and "local_file_paths" in result["data"]:
|
||
for file_path in result["data"]["local_file_paths"]:
|
||
full_path = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), file_path.lstrip('.'))
|
||
if os.path.exists(full_path):
|
||
logger.info(f"本地文件已保存: {full_path}")
|
||
else:
|
||
logger.warning(f"本地文件不存在: {full_path}")
|
||
|
||
return True, result
|
||
else:
|
||
logger.error(f"生成图片失败: HTTP状态码={response.status_code}, 响应内容={response.text}")
|
||
return False, None
|
||
except Exception as e:
|
||
logger.exception(f"生成图片时发生异常: {str(e)}")
|
||
return False, None
|
||
|
||
|
||
def main():
|
||
"""主函数，运行单元测试"""
|
||
logger.info("===== 开始测试QWenImage接口 ====")
|
||
|
||
# 1. 测试获取配置接口
|
||
logger.info("\n1. 测试获取配置接口")
|
||
config_success, config_data = test_get_config()
|
||
|
||
# 2. 测试生成图片接口 - 基本测试
|
||
logger.info("\n2. 测试生成图片接口 - 基本测试")
|
||
basic_prompt = "一只可爱的小猫在草地上玩耍，阳光明媚，高清细节"
|
||
generate_success, generate_data = test_generate_image(
|
||
prompt=basic_prompt,
|
||
n=1,
|
||
size="1328*1328",
|
||
save_local=True
|
||
)
|
||
|
||
# 3. 测试生成图片接口 - 不同参数
|
||
if config_success:
|
||
supported_sizes = config_data["data"].get("supported_sizes", ["1328*1328"])
|
||
max_images = min(2, config_data["data"].get("max_images_per_request", 1)) # 为了测试效率，最多请求2张
|
||
|
||
logger.info(f"\n3. 测试生成图片接口 - 不同参数(size={supported_sizes[0]}, n={max_images})")
|
||
different_prompt = "美丽的山水风景画，中国水墨画风格"
|
||
test_generate_image(
|
||
prompt=different_prompt,
|
||
n=max_images,
|
||
size=supported_sizes[0],
|
||
save_local=True
|
||
)
|
||
|
||
logger.info("\n===== QWenImage接口测试完成 =====")
|
||
|
||
# 输出测试结果摘要
|
||
success_count = sum([config_success, generate_success])
|
||
total_count = 2 # 基本测试的接口数量
|
||
logger.info(f"测试结果: {success_count}/{total_count} 接口测试成功")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main() |