import asyncio import base64 import importlib.util import io import json import sys import threading import types import unittest from pathlib import Path from unittest.mock import AsyncMock, patch from PIL import Image ROOT = Path(__file__).resolve().parents[1] def _load_download_module(): package = types.ModuleType("comfyui_o1key") package.__path__ = [str(ROOT)] utils_package = types.ModuleType("comfyui_o1key.utils") utils_package.__path__ = [str(ROOT / "utils")] clients_package = types.ModuleType("comfyui_o1key.clients") clients_package.__path__ = [str(ROOT / "clients")] sys.modules[package.__name__] = package sys.modules[utils_package.__name__] = utils_package sys.modules[clients_package.__name__] = clients_package http_error_spec = importlib.util.spec_from_file_location( "comfyui_o1key.utils.http_error", ROOT / "utils" / "http_error.py", ) http_error = importlib.util.module_from_spec(http_error_spec) sys.modules[http_error_spec.name] = http_error http_error_spec.loader.exec_module(http_error) gemini_module = types.ModuleType("comfyui_o1key.clients.gemini_client") gemini_module.GeminiAPIClient = object sys.modules[gemini_module.__name__] = gemini_module module_spec = importlib.util.spec_from_file_location( "comfyui_o1key.utils.nano_banana_async", ROOT / "utils" / "nano_banana_async.py", ) module = importlib.util.module_from_spec(module_spec) sys.modules[module_spec.name] = module module_spec.loader.exec_module(module) return module NANO_ASYNC = _load_download_module() def _load_batch_module(): torch_module = types.ModuleType("torch") torch_module.Tensor = type("Tensor", (), {}) sys.modules["torch"] = torch_module numpy_module = types.ModuleType("numpy") numpy_module.random = types.SimpleNamespace(seed=lambda _seed: None) sys.modules["numpy"] = numpy_module comfy_api = types.ModuleType("comfy_api") comfy_latest = types.ModuleType("comfy_api.latest") comfy_latest.io = types.SimpleNamespace(ComfyNode=object, NodeOutput=object) sys.modules["comfy_api"] = comfy_api sys.modules["comfy_api.latest"] = comfy_latest comfy = types.ModuleType("comfy") comfy.__path__ = [] comfy_utils = types.ModuleType("comfy.utils") comfy_utils.ProgressBar = object comfy_model_management = types.ModuleType("comfy.model_management") comfy_model_management.processing_interrupted = lambda: False comfy_model_management.InterruptProcessingException = type( "InterruptProcessingException", (RuntimeError,), {}, ) sys.modules["comfy"] = comfy sys.modules["comfy.utils"] = comfy_utils sys.modules["comfy.model_management"] = comfy_model_management folder_paths = types.ModuleType("folder_paths") folder_paths.get_output_directory = lambda: str(ROOT) sys.modules["folder_paths"] = folder_paths psutil = types.ModuleType("psutil") psutil.Process = object sys.modules["psutil"] = psutil image_utils = types.ModuleType("comfyui_o1key.utils.image_utils") image_utils.tensor_to_pil = lambda _value: [] image_utils.pil_to_tensor = lambda value: value image_utils.parse_batch_prompts = lambda _value: [] sys.modules[image_utils.__name__] = image_utils file_utils = types.ModuleType("comfyui_o1key.utils.file_utils") file_utils.ImageInfo = type("ImageInfo", (), {}) for name in ( "load_images_from_folder", "pair_images_indexed", "pair_images_by_name", "pair_images_cartesian", "generate_timestamp_filename", "save_image", ): setattr(file_utils, name, lambda *_args, **_kwargs: []) sys.modules[file_utils.__name__] = file_utils config = types.ModuleType("comfyui_o1key.utils.config") config.NETWORK_ROUTE_OPTIONS = [] config.get_base_url_by_route = lambda _route: "https://example.invalid" config.get_api_key_or_raise = lambda _name: "test" sys.modules[config.__name__] = config models_config = types.ModuleType("comfyui_o1key.models_config") models_config.get_model_supported_aspect_ratios = lambda _model: [] models_config.get_all_supported_aspect_ratios = lambda: [] models_config.get_model_supported_resolutions = lambda _model: [] models_config.get_all_supported_resolutions = lambda: [] sys.modules[models_config.__name__] = models_config module_spec = importlib.util.spec_from_file_location( "comfyui_o1key.nodes.batch_nano_banana", ROOT / "nodes" / "batch_nano_banana.py", ) module = importlib.util.module_from_spec(module_spec) sys.modules[module_spec.name] = module module_spec.loader.exec_module(module) return module BATCH_NODE = _load_batch_module() def _png_bytes(): buffer = io.BytesIO() Image.new("RGB", (2, 2), (12, 34, 56)).save(buffer, format="PNG") return buffer.getvalue() class _FakeContent: def __init__(self, chunks, delay=0): self._chunks = chunks self._delay = delay async def iter_chunked(self, _chunk_size): for chunk in self._chunks: if self._delay: await asyncio.sleep(self._delay) yield chunk class _FakeResponse: def __init__(self, session, chunks, headers=None, delay=0, status=200): self._session = session self.status = status self.headers = headers or {} self.content = _FakeContent(chunks, delay=delay) async def __aenter__(self): self._session.active += 1 self._session.max_active = max(self._session.max_active, self._session.active) return self async def __aexit__(self, _exc_type, _exc, _tb): self._session.active -= 1 class _FakeSession: def __init__(self, chunks, headers=None, delay=0, status=200): self._chunks = chunks self._headers = headers self._delay = delay self._status = status self.calls = [] self.active = 0 self.max_active = 0 def get(self, url, **_kwargs): self.calls.append(url) return _FakeResponse( self, self._chunks, headers=self._headers, delay=self._delay, status=self._status, ) class DownloadTests(unittest.IsolatedAsyncioTestCase): async def test_successful_task_query_keeps_transport_trace_silent(self): body = json.dumps( {"task_id": "nano-trace", "status": "SUCCESS", "data": {"images": []}}, separators=(",", ":"), ).encode("utf-8") session = _FakeSession( [body], headers={"Content-Length": str(len(body))}, ) with patch("builtins.print") as print_mock: payload = await NANO_ASYNC._poll_task( session, "https://example.invalid", "test-key", "nano-trace", "Nano Banana", log_success=False, initial_delay=False, ) self.assertEqual(payload["status"], "SUCCESS") rendered = " ".join(str(value) for call in print_mock.call_args_list for value in call.args) self.assertNotIn("任务查询传输追踪", rendered) async def test_task_query_retries_when_content_length_is_short(self): body = b'{"task_id":"nano-short","status":"SUCCESS"}' session = _FakeSession( [body], headers={"Content-Length": str(len(body) + 12)}, ) with patch.object( NANO_ASYNC, "_interruptible_sleep", new=AsyncMock(), ): with self.assertRaisesRegex( RuntimeError, rf"Content-Length={len(body) + 12}.*received={len(body)}B.*length_check=mismatch", ): await NANO_ASYNC._poll_task( session, "https://example.invalid", "test-key", "nano-short", "Nano Banana", log_success=False, initial_delay=False, ) self.assertEqual(len(session.calls), 4) def test_unparseable_response_log_never_prints_partial_base64(self): partial_secret = "A" * 4097 with patch("builtins.print") as print_mock: NANO_ASYNC._log_body( "task response", '{"data":{"images":[{"b64_json":"' + partial_secret, ) rendered = " ".join(str(value) for call in print_mock.call_args_list for value in call.args) self.assertNotIn(partial_secret, rendered) self.assertIn("content omitted", rendered) async def test_downloaded_result_is_decoded_only_once(self): image_bytes = _png_bytes() session = _FakeSession([image_bytes]) original_open = NANO_ASYNC._open_result_image with patch.object(NANO_ASYNC, "_open_result_image", wraps=original_open) as open_image: image = await NANO_ASYNC._image_from_url_or_data( "https://example.invalid/result.png", session, ) self.assertEqual(image.size, (2, 2)) self.assertEqual(open_image.call_count, 1) async def test_deduplicates_urls_and_caps_parallel_downloads_at_50(self): image_bytes = _png_bytes() session = _FakeSession( [image_bytes], delay=0.01, ) urls = [f"https://example.invalid/{index}.png" for index in range(60)] payload = { "images": urls, "result": {"image_url": urls[0]}, } images = await NANO_ASYNC._parse_direct_images( payload, session, download_semaphore=asyncio.Semaphore(50), ) self.assertEqual(len(images), 60) self.assertEqual(len(session.calls), 60) self.assertEqual(session.max_active, 50) async def test_retries_when_downloaded_bytes_are_not_a_complete_image(self): session = _FakeSession([b"not a valid image"]) original_delays = NANO_ASYNC._DOWNLOAD_RETRY_DELAYS NANO_ASYNC._DOWNLOAD_RETRY_DELAYS = (0, 0) try: with self.assertRaisesRegex(RuntimeError, "transport failed"): await NANO_ASYNC._image_from_url_or_data( "https://example.invalid/truncated.png", session, ) finally: NANO_ASYNC._DOWNLOAD_RETRY_DELAYS = original_delays self.assertEqual(len(session.calls), 3) async def test_interrupt_check_cancels_during_stream(self): session = _FakeSession([b"1234", b"5678"], delay=0.01) checks = 0 def check_interrupt(): nonlocal checks checks += 1 if checks >= 2: raise asyncio.CancelledError() with self.assertRaises(asyncio.CancelledError): await NANO_ASYNC._download_image_bytes( "https://example.invalid/cancel.png", session, check_interrupt=check_interrupt, ) async def test_refetches_same_task_when_inline_base64_is_incomplete(self): invalid_payload = { "status": "SUCCESS", "data": {"images": [{"b64_json": "truncated-base64"}]}, } valid_payload = { "status": "SUCCESS", "data": {"images": [{ "b64_json": base64.b64encode(_png_bytes()).decode("ascii"), }]}, } session = object() with ( patch.object( NANO_ASYNC, "_poll_task", new=AsyncMock(return_value=valid_payload), ) as poll_task, patch.object( NANO_ASYNC, "_interruptible_sleep", new=AsyncMock(), ) as retry_sleep, ): final_payload, parsed = await NANO_ASYNC._parse_completed_task_images_with_retry( invalid_payload, session=session, base_url="https://example.invalid", api_key="test-key", task_id="task-1", node_label="Nano Banana", ) images, metrics = parsed try: self.assertIs(final_payload, valid_payload) self.assertEqual(len(images), 1) self.assertEqual(images[0].size, (2, 2)) self.assertEqual(metrics["inline_images"], 1) retry_sleep.assert_awaited_once() poll_task.assert_awaited_once_with( session, "https://example.invalid", "test-key", "task-1", "Nano Banana", check_interrupt=None, log_body_enabled=False, progress_callback=None, log_success=False, initial_delay=False, ) finally: for image in images: image.close() class BatchCancellationTests(unittest.IsolatedAsyncioTestCase): def test_timeout_scales_per_50_task_batch(self): per_batch = BATCH_NODE._PER_BATCH_TIMEOUT_SECONDS grace = BATCH_NODE._BATCH_TIMEOUT_GRACE_SECONDS self.assertEqual(BATCH_NODE._batch_timeout_seconds(1), per_batch + grace) self.assertEqual(BATCH_NODE._batch_timeout_seconds(50), per_batch + grace) self.assertEqual(BATCH_NODE._batch_timeout_seconds(51), per_batch * 2 + grace) async def test_stop_event_cancels_active_coroutine(self): stop_event = threading.Event() started = asyncio.Event() cleaned_up = asyncio.Event() async def active_work(): started.set() try: await asyncio.sleep(60) finally: cleaned_up.set() task = asyncio.create_task( BATCH_NODE._run_with_interrupt(active_work(), stop_event) ) await started.wait() stop_event.set() with self.assertRaises(BATCH_NODE._BatchStopRequested): await asyncio.wait_for(task, timeout=1) self.assertTrue(cleaned_up.is_set()) if __name__ == "__main__": unittest.main()