comfyui_o1key/tests/test_nano_banana_async_downloads.py

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()