"""Offline contracts for the native Omni Flash video node.""" import sys import json import tempfile import unittest from pathlib import Path from types import SimpleNamespace from unittest.mock import AsyncMock, patch PLUGIN_ROOT = Path(__file__).resolve().parents[1] sys.path.insert(0, str(PLUGIN_ROOT.parent)) import comfy.options # noqa: E402 comfy.options.enable_args_parsing() sys.argv = [sys.argv[0], "--cpu"] from comfyui_o1key.clients.omni_flash_client import OmniFlashClient, build_video_body, _log_response_body, _response_error, _safe_error, _submission_payload, _task_error # noqa: E402 from comfyui_o1key.nodes.omni_flash_video import O1keyOmniFlashVideo, _make_progress_callback # noqa: E402 sys.argv = [sys.argv[0]] def inputs(**changes): result = {"提示词": "clouds over mountains", "生成模式": "文生视频", "分辨率": "720p", "宽高比": "16:9"} result.update(changes) return result class OmniFlashContractTests(unittest.TestCase): def test_client_uses_existing_o1key_token_and_route(self): with patch("comfyui_o1key.clients.omni_flash_client.get_api_key_or_raise", return_value="fake-key") as read_key, \ patch("comfyui_o1key.clients.omni_flash_client.get_base_url_by_route", return_value="https://api.o1key.invalid"): client = OmniFlashClient() read_key.assert_called_once_with("O1KEY_API_KEY") self.assertEqual(client.base_url, "https://api.o1key.invalid") def test_schema_is_single_native_video_node(self): schema = O1keyOmniFlashVideo.define_schema() schema.validate() self.assertEqual(schema.node_id, "O1keyOmniFlashVideo") self.assertNotIn("模型", [item.id for item in schema.inputs]) self.assertEqual([item.id for item in schema.outputs], ["VIDEO"]) self.assertTrue(schema.is_output_node) self.assertTrue(schema.not_idempotent) def test_repeated_runs_get_distinct_cache_fingerprints(self): first = O1keyOmniFlashVideo.fingerprint_inputs(**inputs()) second = O1keyOmniFlashVideo.fingerprint_inputs(**inputs()) self.assertNotEqual(first, second) def test_documented_payload_modes(self): base = dict(prompt="sky", resolution="720p", aspect_ratio="16:9") text = build_video_body(model="omni_flash_8s", mode="text", **base) self.assertNotIn("input_reference", text) url = "https://example.invalid/first.png" reference = build_video_body(model="omni_flash_8s", mode="reference", references=[url], **base) self.assertEqual(reference["input_reference"], url) frames = build_video_body(model="omni_flash_10s", mode="first_last_frame", references=[url, url], **base) self.assertTrue(frames["first_last_frame"]) self.assertEqual(frames["input_reference"], [url, url]) first_only = build_video_body(model="omni_flash_10s", mode="first_last_frame", references=[url], **base) self.assertEqual(first_only["input_reference"], url) self.assertNotIn("first_last_frame", first_only) edit = build_video_body(model="omni_flash_abra_edit", mode="edit", source_video_url="https://example.invalid/video.mp4", **base) self.assertEqual(edit["model"], "omni_flash_abra_edit") def test_invalid_payload_and_sensitive_error(self): with self.assertRaises(ValueError): build_video_body(model="omni_flash_8s", mode="reference", prompt="sky", resolution="720p", aspect_ratio="16:9") self.assertNotIn("token=secret", _safe_error("https://host.invalid/a?token=secret")) self.assertNotIn("token123", _safe_error("Bearer token123")) def test_gateway_400_preserves_useful_detail_without_urls(self): message = _response_error({"message": "invalid image https://example.invalid/a?token=secret"}, 400) self.assertIn("invalid image", message) self.assertNotIn("token=secret", message) self.assertIn("HTTP 400", message) def test_documented_error_codes_have_specific_messages(self): self.assertIn("令牌管理", _response_error({"error": {"code": "invalid_api_key"}}, 401)) self.assertIn("余额不足", _response_error({"error": {"code": "insufficient_balance"}}, 402)) self.assertIn("请求过于频繁", _response_error({"error": {"code": "rate_limit_exceeded"}}, 429)) self.assertIn("参考图", _task_error({"data": {"error": {"code": "image_url_required_for_i2v"}}})) def test_response_log_keeps_status_and_masks_sensitive_values(self): body = json.dumps({ "status": "completed", "url": "https://cdn.example.invalid/video?token=secret", "api_key": "secret-key", "result": {"progress": 100}, }) with patch("builtins.print") as output: _log_response_body("查询", 200, body) logged = output.call_args.args[0] self.assertIn('"status":"completed"', logged) self.assertIn('"progress":100', logged) self.assertNotIn("secret-key", logged) self.assertNotIn("token=secret", logged) def test_several_references_use_repeated_multipart_field(self): first, last = "https://example.invalid/first.png", "https://example.invalid/last.png" body = build_video_body(model="omni_flash_8s", prompt="sky", resolution="720p", aspect_ratio="16:9", mode="first_last_frame", references=[first, last]) submitted = _submission_payload(body) self.assertNotIn("json", submitted) form = submitted["data"] self.assertTrue(form.is_multipart) fields = [(options["name"], value) for options, _headers, value in form._fields] self.assertEqual([value for name, value in fields if name == "input_reference"], [first, last]) self.assertIn(("first_last_frame", "true"), fields) def test_first_last_still_requires_first_frame(self): with self.assertRaisesRegex(ValueError, "首帧"): build_video_body(model="omni_flash_8s", mode="first_last_frame", prompt="sky", resolution="720p", aspect_ratio="16:9", references=[]) def test_edit_source_limit_before_upload(self): with tempfile.TemporaryDirectory() as directory: path = Path(directory) / "source.mp4" with path.open("wb") as stream: stream.truncate(20 * 1024 * 1024 + 1) with self.assertRaisesRegex(ValueError, "20 MB"): O1keyOmniFlashVideo._check_source_video(SimpleNamespace(get_stream_source=lambda: str(path))) def test_provider_progress_maps_to_native_node_bar(self): with patch("comfy.utils.ProgressBar") as progress_bar: update = _make_progress_callback() update("polling", 0, "task_12345678") update("polling", 50, "task_12345678") update("polling", 25, "task_12345678") update("downloading", 100, "task_12345678") update("done", 100, "task_12345678") self.assertEqual( [call.args for call in progress_bar.return_value.update_absolute.call_args_list], [(0, 100), (50, 100), (99, 100), (100, 100)], ) class OmniFlashExecutionTests(unittest.IsolatedAsyncioTestCase): async def test_invalid_mode_rejects_before_upload_and_submit(self): with patch("comfyui_o1key.nodes.omni_flash_video.upload_image", new_callable=AsyncMock) as upload, \ patch("comfyui_o1key.nodes.omni_flash_video.OmniFlashClient") as client: with self.assertRaises(ValueError): await O1keyOmniFlashVideo.execute(**inputs(**{"生成模式": "参考图视频"})) upload.assert_not_awaited() client.assert_not_called() async def test_native_queue_execution_returns_video(self): async def fake_generate(_client, body, path, progress=None): self.assertEqual(body["model"], "omni_flash_10s") self.assertNotIn("input_reference", body) self.assertEqual(_client.api_key, "fake-key") self.assertEqual(_client.base_url, "https://api.o1key.invalid") Path(path).write_bytes(b"video") return "task_12345678" with tempfile.TemporaryDirectory() as directory: with patch("comfyui_o1key.nodes.omni_flash_video.folder_paths.get_output_directory", return_value=directory), \ patch("comfyui_o1key.nodes.omni_flash_video.get_api_key_or_raise", return_value="fake-key"), \ patch("comfyui_o1key.nodes.omni_flash_video.get_base_url_by_route", return_value="https://api.o1key.invalid"), \ patch("comfyui_o1key.nodes.omni_flash_video.OmniFlashClient.generate", fake_generate), \ patch("comfyui_o1key.nodes.omni_flash_video.InputImpl.VideoFromFile", side_effect=lambda path: path): result = await O1keyOmniFlashVideo.execute(**inputs()) self.assertEqual(len(list(Path(directory, "omni_flash").glob("*.mp4"))), 1) self.assertFalse(list(Path(directory, "omni_flash").glob("*.part"))) self.assertIsNone(result.ui) async def test_only_first_frame_uploads_one_image(self): async def fake_generate(_client, body, path, progress=None): self.assertNotIn("first_last_frame", body) self.assertEqual(body["input_reference"], "https://example.invalid/first.png") Path(path).write_bytes(b"video") image = SimpleNamespace(convert=lambda _mode: "rgb-image") with tempfile.TemporaryDirectory() as directory: with patch("comfyui_o1key.nodes.omni_flash_video.folder_paths.get_output_directory", return_value=directory), \ patch("comfyui_o1key.nodes.omni_flash_video.get_api_key_or_raise", return_value="fake-key"), \ patch("comfyui_o1key.nodes.omni_flash_video.get_base_url_by_route", return_value="https://api.o1key.invalid"), \ patch("comfyui_o1key.nodes.omni_flash_video.tensor_to_pil", return_value=[image]), \ patch("comfyui_o1key.nodes.omni_flash_video.upload_image", new_callable=AsyncMock, return_value="https://example.invalid/first.png") as upload, \ patch("comfyui_o1key.nodes.omni_flash_video.OmniFlashClient.generate", fake_generate), \ patch("comfyui_o1key.nodes.omni_flash_video.InputImpl.VideoFromFile", side_effect=lambda path: path): await O1keyOmniFlashVideo.execute(**inputs(**{"生成模式": "首尾帧", "首帧图片": object()})) upload.assert_awaited_once() async def test_client_submits_once_then_polls_and_downloads(self): class Response: def __init__(self, data, content_type="application/json"): self.data, self.status = data, 200 self.headers = {"Content-Type": content_type} async def __aenter__(self): return self async def __aexit__(self, *_): return False async def json(self, **_): return self.data async def text(self): return json.dumps(self.data) class Session: def __init__(self): self.posts = [] self.gets = [] self.statuses = [ Response({"status": "unknown", "progress": "25%"}), Response({"status": "success", "data": {"task_status": "processing", "progress": "50%"}}), Response({"data": {"status": "succeeded"}}), Response({}, "video/mp4"), ] async def __aenter__(self): return self async def __aexit__(self, *_): return False def post(self, url, **kwargs): self.posts.append((url, kwargs)) return Response({"id": "task_12345678"}) def get(self, *_args, **kwargs): self.gets.append(kwargs) return self.statuses.pop(0) session = Session() client = OmniFlashClient(base_url="https://example.invalid", api_key="fake-key") body = build_video_body(model="omni_flash_8s", prompt="sky", resolution="720p", aspect_ratio="16:9", mode="text") with patch("comfyui_o1key.clients.omni_flash_client.aiohttp.ClientSession", return_value=session), \ patch("comfyui_o1key.clients.omni_flash_client.interruptible_sleep", new_callable=AsyncMock), \ patch("comfyui_o1key.clients.omni_flash_client.download_video_to_file", new_callable=AsyncMock) as download: task_id = await client.generate(body, "result.mp4") self.assertEqual(task_id, "task_12345678") self.assertEqual(len(session.posts), 1) self.assertEqual(session.posts[0][1]["json"], body) self.assertNotIn("X-No-Watermark", session.posts[0][1]["headers"]) self.assertIn("/v1/videos/task_12345678/content", download.call_args.args[1]) edit_body = build_video_body( model="omni_flash_abra_edit", prompt="edit", resolution="720p", aspect_ratio="16:9", mode="edit", source_video_url="https://example.invalid/source.mp4", ) edit_session = Session() with patch("comfyui_o1key.clients.omni_flash_client.aiohttp.ClientSession", return_value=edit_session), \ patch("comfyui_o1key.clients.omni_flash_client.interruptible_sleep", new_callable=AsyncMock), \ patch("comfyui_o1key.clients.omni_flash_client.download_video_to_file", new_callable=AsyncMock): await client.generate(edit_body, "edited.mp4") self.assertEqual(edit_session.posts[0][1]["headers"]["X-No-Watermark"], "video") self.assertTrue(all("X-No-Watermark" not in request["headers"] for request in edit_session.gets)) async def test_completed_content_json_uses_result_link_without_auth_header(self): class Response: status = 200 headers = {"Content-Type": "application/json"} async def __aenter__(self): return self async def __aexit__(self, *_): return False async def json(self, **_): return {"data": {"download_url": "https://cdn.example.invalid/video.mp4"}} async def text(self): return json.dumps(await self.json()) session = SimpleNamespace(get=lambda *_args, **_kwargs: Response()) client = OmniFlashClient(base_url="https://example.invalid", api_key="fake-key") with patch("comfyui_o1key.clients.omni_flash_client.download_video_to_file", new_callable=AsyncMock) as download: await client._download_completed(session, "task_12345678", "result.mp4", {"Authorization": "Bearer fake-key"}, {"status": "completed"}) self.assertEqual(download.call_args.args[1], "https://cdn.example.invalid/video.mp4") self.assertIsNone(download.call_args.kwargs["headers"]) if __name__ == "__main__": unittest.main()