#!/usr/bin/env python3
"""
即梦AI图像生成服务测试客户端
用于测试 MCP 服务的功能
"""
import os
import sys
def test_with_service_directly():
"""
直接测试 service.py（不通过 MCP）
这是最简单的测试方式
"""
from service import create_service
# 从环境变量获取 token
token = os.getenv("JIMENG_TOKEN", "bd840c796b50b4f711817b91288f9507")
if not token:
print("错误: 请设置 JIMENG_TOKEN 环境变量")
print("使用方法: export JIMENG_TOKEN='your_token_here'")
sys.exit(1)
# 创建服务实例
print("正在创建服务实例...")
service = create_service(token)
# 测试参数
prompt = "秋日街道旁的时尚店铺，一位长发微卷的亚洲女性站在弧形玻璃橱窗边"
save_folder = "./test_output"
width = 2304
height = 1728
print(f"\n开始测试图像生成:")
print(f"提示词: {prompt[:50]}...")
print(f"保存路径: {save_folder}")
print(f"图像尺寸: {width}x{height}")
print(f"\n正在提交请求...")
# 调用生成图像
result = service.generate_image(
prompt=prompt,
save_folder=save_folder,
width=width,
height=height
)
# 打印结果
print("\n" + "="*60)
if result.get("success"):
print("✅ 图像生成成功!")
print(f"\n提交ID: {result.get('submit_id')}")
print(f"图片总数: {result.get('total_images')}")
print(f"保存文件夹: {result.get('save_folder')}")
print("\n图片URL列表:")
for idx, url in enumerate(result.get('image_urls', []), 1):
print(f" {idx}. {url}")
print("\n已保存文件:")
for file_path in result.get('saved_files', []):
print(f" ✓ {file_path}")
if result.get('failed_downloads'):
print("\n下载失败:")
for fail in result['failed_downloads']:
print(f" ✗ 第{fail['index']}张: {fail['error']}")
else:
print("❌ 图像生成失败!")
print(f"错误信息: {result.get('error')}")
print("="*60)
return result
def test_with_mcp_client():
"""
使用 MCP 客户端测试（需要安装 mcp 包）
"""
try:
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
except ImportError:
print("错误: 需要安装 mcp 包")
print("运行: pip install mcp")
sys.exit(1)
print("MCP 客户端测试功能尚未实现")
print("请使用直接测试模式: python test_client.py --direct")
def show_help():
"""显示帮助信息"""
help_text = """
即梦AI图像生成服务测试工具
使用方法:
python test_client.py [选项]
选项:
--direct 直接测试 service.py（推荐）
--mcp 使用 MCP 客户端测试
--help, -h 显示此帮助信息
环境变量:
JIMENG_TOKEN 即梦平台的认证token（必需）
示例:
# 1. 设置环境变量
export JIMENG_TOKEN="your_token_here"
# 2. 运行直接测试
python test_client.py --direct
# 3. 查看生成的图片
ls -lh ./test_output/
注意:
- 确保已安装依赖: pip install -r requirements.txt
- 首次运行会在当前目录创建 ./test_output 文件夹
- 图像生成需要等待一段时间（通常30秒-3分钟）
"""
print(help_text)
if __name__ == "__main__":
# 解析命令行参数
args = sys.argv[1:]
if not args or "--direct" in args:
# 默认使用直接测试
test_with_service_directly()
elif "--mcp" in args:
test_with_mcp_client()
elif "--help" in args or "-h" in args:
show_help()
else:
print(f"未知选项: {args[0]}")
print("使用 --help 查看帮助")
sys.exit(1)
