249 lines
8.3 KiB
Python
249 lines
8.3 KiB
Python
"""
|
||
本地测试脚本 - 用于测试 Sora 自定义节点
|
||
|
||
使用方法：
|
||
1. 直接运行测试（需要 mock API）：
|
||
python test_sora_node.py
|
||
|
||
2. 测试节点类的基本功能：
|
||
python test_sora_node.py --test-class
|
||
|
||
3. 测试 API 客户端（需要真实的 API key）：
|
||
python test_sora_node.py --test-client --api-key YOUR_KEY
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
import sys
|
||
import unittest
|
||
from unittest.mock import Mock, patch, MagicMock
|
||
from urllib.error import HTTPError
|
||
|
||
# 导入节点代码
|
||
from sora_node import SoraClient, SoraTextToVideo, SoraAPIError, DEFAULT_SORA_BASE_URL
|
||
|
||
|
||
class TestSoraClient(unittest.TestCase):
|
||
"""测试 SoraClient 类"""
|
||
|
||
def test_client_init_with_api_key(self):
|
||
"""测试客户端初始化"""
|
||
client = SoraClient(api_key="test-key")
|
||
self.assertEqual(client.api_key, "test-key")
|
||
self.assertEqual(client.base_url, DEFAULT_SORA_BASE_URL)
|
||
|
||
def test_client_init_without_api_key(self):
|
||
"""测试缺少 API key 时抛出错误"""
|
||
with self.assertRaises(ValueError):
|
||
SoraClient(api_key="")
|
||
|
||
@patch('sora_node.urllib.request.urlopen')
|
||
def test_create_job_success(self, mock_urlopen):
|
||
"""测试创建任务成功"""
|
||
# Mock HTTP 响应
|
||
mock_response = MagicMock()
|
||
mock_response.read.return_value = b'{"id": "job-123", "status": "submitted"}'
|
||
mock_urlopen.return_value.__enter__.return_value = mock_response
|
||
|
||
client = SoraClient(api_key="test-key")
|
||
result = client.create_job({"prompt": "test prompt"})
|
||
|
||
self.assertEqual(result["id"], "job-123")
|
||
self.assertEqual(result["status"], "submitted")
|
||
|
||
@patch('sora_node.urllib.request.urlopen')
|
||
def test_create_job_error(self, mock_urlopen):
|
||
"""测试创建任务失败"""
|
||
# Mock HTTP 错误
|
||
mock_error = HTTPError("http://test.com", 400, "Bad Request", {}, None)
|
||
mock_error.read.return_value = b'{"error": "Invalid prompt"}'
|
||
mock_urlopen.side_effect = mock_error
|
||
|
||
client = SoraClient(api_key="test-key")
|
||
with self.assertRaises(SoraAPIError):
|
||
client.create_job({"prompt": "test"})
|
||
|
||
|
||
class TestSoraTextToVideo(unittest.TestCase):
|
||
"""测试 SoraTextToVideo 节点类"""
|
||
|
||
def test_input_types(self):
|
||
"""测试节点输入类型定义"""
|
||
node = SoraTextToVideo()
|
||
input_types = node.INPUT_TYPES()
|
||
|
||
self.assertIn("required", input_types)
|
||
self.assertIn("optional", input_types)
|
||
self.assertIn("prompt", input_types["required"])
|
||
|
||
def test_resolve_api_key_from_env(self):
|
||
"""测试从环境变量解析 API key"""
|
||
# 测试 SORA_API_KEY
|
||
os.environ["SORA_API_KEY"] = "env-key-123"
|
||
from sora_node import _resolve_api_key
|
||
key = _resolve_api_key(None)
|
||
self.assertEqual(key, "env-key-123")
|
||
|
||
# 测试 OPENAI_API_KEY（fallback）
|
||
del os.environ["SORA_API_KEY"]
|
||
os.environ["OPENAI_API_KEY"] = "openai-key-456"
|
||
key = _resolve_api_key(None)
|
||
self.assertEqual(key, "openai-key-456")
|
||
|
||
# 清理
|
||
if "OPENAI_API_KEY" in os.environ:
|
||
del os.environ["OPENAI_API_KEY"]
|
||
|
||
@patch('sora_node.SoraClient')
|
||
def test_generate_basic(self, mock_client_class):
|
||
"""测试基本的生成功能"""
|
||
# Mock 客户端
|
||
mock_client = MagicMock()
|
||
mock_client_class.return_value = mock_client
|
||
|
||
# Mock API 响应
|
||
mock_client.create_job.return_value = {
|
||
"id": "test-job-123",
|
||
"status": "submitted"
|
||
}
|
||
|
||
node = SoraTextToVideo()
|
||
job_id, status, result = node.generate(
|
||
prompt="test prompt",
|
||
api_key="test-key",
|
||
wait_for_result=False
|
||
)
|
||
|
||
self.assertEqual(job_id, "test-job-123")
|
||
self.assertEqual(status, "submitted")
|
||
result_data = json.loads(result)
|
||
self.assertEqual(result_data["job_id"], "test-job-123")
|
||
|
||
@patch('sora_node.SoraClient')
|
||
def test_generate_with_polling(self, mock_client_class):
|
||
"""测试带轮询的生成"""
|
||
mock_client = MagicMock()
|
||
mock_client_class.return_value = mock_client
|
||
|
||
# Mock 创建任务
|
||
mock_client.create_job.return_value = {
|
||
"id": "test-job-456",
|
||
"status": "processing"
|
||
}
|
||
|
||
# Mock 轮询响应
|
||
mock_client.fetch_job.return_value = {
|
||
"id": "test-job-456",
|
||
"status": "succeeded",
|
||
"result_url": "https://example.com/video.mp4"
|
||
}
|
||
|
||
node = SoraTextToVideo()
|
||
job_id, status, result = node.generate(
|
||
prompt="test prompt",
|
||
api_key="test-key",
|
||
wait_for_result=True,
|
||
max_wait_seconds=10.0,
|
||
poll_interval=1.0
|
||
)
|
||
|
||
self.assertEqual(job_id, "test-job-456")
|
||
self.assertEqual(status, "succeeded")
|
||
# 验证轮询被调用
|
||
mock_client.fetch_job.assert_called()
|
||
|
||
def test_extract_video_url(self):
|
||
"""测试视频 URL 提取"""
|
||
node = SoraTextToVideo()
|
||
|
||
# 测试 result_url
|
||
response1 = {"result_url": "https://example.com/video1.mp4"}
|
||
url1 = node._extract_video_url(response1)
|
||
self.assertEqual(url1, "https://example.com/video1.mp4")
|
||
|
||
# 测试 outputs 数组
|
||
response2 = {
|
||
"outputs": [{"url": "https://example.com/video2.mp4"}]
|
||
}
|
||
url2 = node._extract_video_url(response2)
|
||
self.assertEqual(url2, "https://example.com/video2.mp4")
|
||
|
||
# 测试空响应
|
||
response3 = {}
|
||
url3 = node._extract_video_url(response3)
|
||
self.assertIsNone(url3)
|
||
|
||
|
||
def test_node_mappings():
|
||
"""测试节点映射是否正确"""
|
||
from sora_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||
|
||
assert "SoraTextToVideo" in NODE_CLASS_MAPPINGS
|
||
assert NODE_CLASS_MAPPINGS["SoraTextToVideo"] == SoraTextToVideo
|
||
assert "SoraTextToVideo" in NODE_DISPLAY_NAME_MAPPINGS
|
||
print("✓ 节点映射测试通过")
|
||
|
||
|
||
def run_basic_tests():
|
||
"""运行基本功能测试（不需要真实 API）"""
|
||
print("=" * 60)
|
||
print("运行基本功能测试...")
|
||
print("=" * 60)
|
||
|
||
# 测试节点映射
|
||
test_node_mappings()
|
||
|
||
# 测试输入类型
|
||
node = SoraTextToVideo()
|
||
input_types = node.INPUT_TYPES()
|
||
assert "required" in input_types
|
||
assert "prompt" in input_types["required"]
|
||
print("✓ 节点输入类型测试通过")
|
||
|
||
# 测试 API key 解析
|
||
os.environ["SORA_API_KEY"] = "test-key-123"
|
||
from sora_node import _resolve_api_key
|
||
key = _resolve_api_key(None)
|
||
assert key == "test-key-123"
|
||
del os.environ["SORA_API_KEY"]
|
||
print("✓ API key 解析测试通过")
|
||
|
||
print("\n所有基本测试通过！")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
|
||
parser = argparse.ArgumentParser(description="测试 Sora 自定义节点")
|
||
parser.add_argument("--test-class", action="store_true", help="运行类测试")
|
||
parser.add_argument("--test-client", action="store_true", help="测试客户端（需要真实 API key）")
|
||
parser.add_argument("--api-key", type=str, help="API key（用于客户端测试）")
|
||
parser.add_argument("--unittest", action="store_true", help="运行完整的 unittest 测试套件")
|
||
|
||
args = parser.parse_args()
|
||
|
||
if args.unittest:
|
||
# 运行完整的 unittest 测试
|
||
unittest.main(argv=[''], exit=False, verbosity=2)
|
||
elif args.test_class:
|
||
# 运行基本类测试
|
||
run_basic_tests()
|
||
elif args.test_client and args.api_key:
|
||
# 测试真实 API 客户端（需要网络连接和有效 API key）
|
||
print("测试真实 API 客户端...")
|
||
print("注意：这需要有效的 API key 和网络连接")
|
||
try:
|
||
client = SoraClient(api_key=args.api_key)
|
||
# 这里可以添加真实的 API 调用测试
|
||
print("客户端初始化成功")
|
||
except Exception as e:
|
||
print(f"客户端测试失败: {e}")
|
||
else:
|
||
# 默认运行基本测试
|
||
run_basic_tests()
|
||
print("\n提示：")
|
||
print(" - 运行完整测试: python test_sora_node.py --unittest")
|
||
print(" - 运行基本测试: python test_sora_node.py --test-class")
|
||
print(" - 测试真实 API: python test_sora_node.py --test-client --api-key YOUR_KEY")
|
||
|