188 lines
5.8 KiB
Python
188 lines
5.8 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
|
||
"""
|
||
百度图像搜索API测试脚本
|
||
简单测试API功能，打印结果
|
||
支持实际API和模拟模式
|
||
"""
|
||
|
||
import os
|
||
import json
|
||
import argparse
|
||
from dotenv import load_dotenv
|
||
from app.api.baidu_image_search import BaiduImageSearch
|
||
from app.api.image_utils import ImageUtils
|
||
|
||
# 导入模拟类
|
||
from mock_baidu_api import MockBaiduImageSearch
|
||
|
||
# 全局变量，用于存储测试过程中的图片ID
|
||
test_cont_sign = None
|
||
|
||
def parse_arguments():
|
||
"""解析命令行参数"""
|
||
parser = argparse.ArgumentParser(description='百度图像搜索API测试脚本')
|
||
parser.add_argument('--mock', action='store_true', help='使用模拟模式，不连接实际API')
|
||
return parser.parse_args()
|
||
|
||
def check_api_keys():
|
||
"""检查API密钥是否已配置"""
|
||
api_key = os.getenv('BAIDU_API_KEY')
|
||
secret_key = os.getenv('BAIDU_SECRET_KEY')
|
||
|
||
if not api_key or not secret_key:
|
||
print("\n\033[91m错误: API密钥未配置\033[0m")
|
||
print("\n请在.env文件中添加以下内容:")
|
||
print("BAIDU_API_KEY=你的API密钥")
|
||
print("BAIDU_SECRET_KEY=你的密钥")
|
||
print("\n您可以从百度AI开放平台获取密钥: https://ai.baidu.com/")
|
||
print("\n或者使用 --mock 参数运行模拟模式: python test_baidu_api.py --mock")
|
||
return False
|
||
|
||
print(f"API密钥: {api_key[:4]}...")
|
||
print(f"Secret密钥: {secret_key[:4]}...")
|
||
return True
|
||
|
||
def main():
|
||
# 解析命令行参数
|
||
args = parse_arguments()
|
||
|
||
# 加载环境变量
|
||
print("加载环境变量...")
|
||
load_dotenv()
|
||
|
||
# 判断是否使用模拟模式
|
||
if args.mock:
|
||
print("\n\033[93m注意: 使用模拟模式，不连接实际API\033[0m")
|
||
api = MockBaiduImageSearch()
|
||
|
||
# 测试添加图片
|
||
test_add_image(api)
|
||
|
||
# 测试搜索图片
|
||
test_search_image(api)
|
||
|
||
# 测试更新图片
|
||
test_update_image(api)
|
||
|
||
# 测试删除图片
|
||
test_delete_image(api)
|
||
return
|
||
|
||
# 如果不是模拟模式，检查API密钥
|
||
if not check_api_keys():
|
||
return
|
||
|
||
# 初始化实际API
|
||
print("\n初始化百度图像搜索API...")
|
||
try:
|
||
api = BaiduImageSearch()
|
||
if api.access_token:
|
||
print(f"获取到的Access Token: {api.access_token[:10]}...（已截断）")
|
||
|
||
# 测试添加图片
|
||
test_add_image(api)
|
||
|
||
# 测试搜索图片
|
||
test_search_image(api)
|
||
|
||
# 测试更新图片
|
||
test_update_image(api)
|
||
|
||
# 测试删除图片
|
||
test_delete_image(api)
|
||
else:
|
||
print("\033[91m错误: 无法获取Access Token\033[0m")
|
||
print("\n尝试使用模拟模式: python test_baidu_api.py --mock")
|
||
except Exception as e:
|
||
print(f"\033[91m错误: {e}\033[0m")
|
||
print("\n尝试使用模拟模式: python test_baidu_api.py --mock")
|
||
|
||
def test_add_image(api):
|
||
"""测试添加图片功能"""
|
||
print("\n===== 测试添加图片 =====")
|
||
|
||
# 使用URL方式添加图片
|
||
test_image_url = "https://www.baidu.com/img/PCtm_d9c8750bed0b3c7d089fa7d55720d6cf.png" # 百度logo
|
||
brief = json.dumps({"name": "测试图片", "id": "test001"}, ensure_ascii=False)
|
||
tags = "测试,图片"
|
||
|
||
try:
|
||
result = api.add_image(
|
||
url=test_image_url,
|
||
brief=brief,
|
||
tags=tags
|
||
)
|
||
print("添加图片结果:")
|
||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||
|
||
# 保存contSign用于后续测试
|
||
if result.get("cont_sign"):
|
||
global test_cont_sign
|
||
test_cont_sign = result.get("cont_sign")
|
||
print(f"已保存图片ID: {test_cont_sign} 用于后续测试")
|
||
except Exception as e:
|
||
print(f"添加图片失败: {e}")
|
||
|
||
def test_search_image(api):
|
||
"""测试搜索图片功能"""
|
||
print("\n===== 测试搜索图片 =====")
|
||
|
||
# 使用URL方式搜索图片
|
||
test_image_url = "https://www.baidu.com/img/PCtm_d9c8750bed0b3c7d089fa7d55720d6cf.png" # 百度logo
|
||
|
||
try:
|
||
result = api.search_image(
|
||
url=test_image_url,
|
||
tags="测试"
|
||
)
|
||
print("搜索图片结果:")
|
||
print(f"找到 {len(result.get('result', []))} 个匹配项")
|
||
|
||
# 只打印前2个结果
|
||
for i, item in enumerate(result.get("result", [])[:2]):
|
||
print(f"结果 {i+1}:")
|
||
print(f" 相似度: {item.get('score', 0)}")
|
||
print(f" 简介: {item.get('brief', '')}")
|
||
except Exception as e:
|
||
print(f"搜索图片失败: {e}")
|
||
|
||
def test_update_image(api):
|
||
"""测试更新图片功能"""
|
||
print("\n===== 测试更新图片 =====")
|
||
|
||
global test_cont_sign
|
||
if not test_cont_sign:
|
||
print("没有可用的图片ID，跳过更新图片测试")
|
||
return
|
||
|
||
try:
|
||
new_brief = json.dumps({"name": "更新后的测试图片", "id": "test001"}, ensure_ascii=False)
|
||
result = api.update_image(test_cont_sign, new_brief)
|
||
print("更新图片结果:")
|
||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||
except Exception as e:
|
||
print(f"更新图片失败: {e}")
|
||
|
||
def test_delete_image(api):
|
||
"""测试删除图片功能"""
|
||
print("\n===== 测试删除图片 =====")
|
||
|
||
global test_cont_sign
|
||
if not test_cont_sign:
|
||
print("没有可用的图片ID，跳过删除图片测试")
|
||
return
|
||
|
||
try:
|
||
result = api.delete_image(test_cont_sign)
|
||
print("删除图片结果:")
|
||
print(json.dumps(result, indent=2, ensure_ascii=False))
|
||
# 清除已删除的图片ID
|
||
test_cont_sign = None
|
||
except Exception as e:
|
||
print(f"删除图片失败: {e}")
|
||
|
||
if __name__ == "__main__":
|
||
main()
|