dsProject/dsLightRag/Test/TestQWen3Image.py

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,  size='1328*1328', save_local=True):
    """测试生成图片接口
    
    Args:
        prompt: 图片描述提示词
        size: 图片尺寸
        save_local: 是否保存到本地
        
    Returns:
        tuple: (是否成功, 结果数据)
    """
    try:
        url = f"{base_url}/generate"
        headers = {"Content-Type": "application/json"}
        data = {
            "prompt": prompt,
            "size": size,
            "save_local": save_local
        }
        
        logger.info(f"调用生成图片接口: {url}")
        logger.info(f"请求参数: prompt={prompt[:50]}...,  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,
        size="1328*1328",
        save_local=True
    )
    
    # 3. 测试生成图片接口 - 不同参数
    if config_success:
        supported_sizes = config_data["data"].get("supported_sizes", ["1328*1328"])

        logger.info(f"\n3. 测试生成图片接口 - 不同参数(size={supported_sizes[0]})")
        different_prompt = "美丽的山水风景画，中国水墨画风格"
        test_generate_image(
            prompt=different_prompt,
            size=supported_sizes[0],
            save_local=True
        )
    
    logger.info("\n===== QWenImage接口测试完成 =====")
    
    # 输出测试结果摘要
    success_count = sum([config_success, generate_success])
    logger.info(f"测试结果: {success_count} 接口测试成功")


if __name__ == "__main__":
    main()