""" 本地测试脚本 - 用于测试 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")