comfyui_o1key/tests/test_nano_banana_temp_uploads.py

import asyncio
import base64
import importlib.util
import io
import json
import random
import sys
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_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_module()


class _FakeResponse:
    def __init__(self, status, payload, headers=None):
        self.status = status
        self._text = json.dumps(payload)
        self.headers = headers or {}

    async def __aenter__(self):
        return self

    async def __aexit__(self, _exc_type, _exc, _tb):
        return False

    async def text(self):
        return self._text


class _FakeSession:
    def __init__(self, response):
        self.response = response
        self.calls = []

    def post(self, url, **kwargs):
        self.calls.append((url, kwargs))
        return self.response


def _noise_image(width=64, height=64):
    rng = random.Random(7)
    data = bytes(rng.randrange(256) for _ in range(width * height * 3))
    return Image.frombytes("RGB", (width, height), data)


def _jpeg_image_bytes(image, quality=95):
    buffer = io.BytesIO()
    image.save(buffer, format="JPEG", quality=quality)
    return buffer.getvalue()


class TempUploadTests(unittest.IsolatedAsyncioTestCase):
    def test_png_upload_payload_keeps_original_bytes(self):
        image = _noise_image(16, 12)
        original = NANO_ASYNC._png_bytes(image)
        image._o1key_original_format = "PNG"
        image._o1key_original_bytes = original

        data, extension, content_type = NANO_ASYNC._image_to_upload_payload(image)

        self.assertEqual(data, original)
        self.assertEqual(extension, ".png")
        self.assertEqual(content_type, "image/png")

    def test_jpeg_upload_payload_keeps_original_bytes(self):
        image = _noise_image(32, 24)
        original = _jpeg_image_bytes(image, quality=95)
        image._o1key_original_format = "JPEG"
        image._o1key_original_bytes = original

        data, extension, content_type = NANO_ASYNC._image_to_upload_payload(image)

        self.assertEqual(data, original)
        self.assertEqual(extension, ".jpg")
        self.assertEqual(content_type, "image/jpeg")

    async def test_upload_posts_multipart_file_and_returns_https_url(self):
        expected_url = "https://cf-api.o1key.com/tmp/input/reference.png"
        response = _FakeResponse(200, {
            "url": expected_url,
            "filename": "reference.png",
            "content_type": "image/png",
            "size": 100,
            "expires_at": 1_787_495_062,
        })
        session = _FakeSession(response)

        urls = await NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
            session=session,
            base_url="https://cf-api.o1key.com/",
            api_key="secret",
            images=[Image.new("RGB", (2, 2), "red")],
        )

        self.assertEqual(urls, [expected_url])
        self.assertEqual(len(session.calls), 1)
        url, kwargs = session.calls[0]
        self.assertEqual(url, "https://cf-api.o1key.com/v1/o1key/uploads")
        self.assertEqual(kwargs["headers"], {"Authorization": "Bearer secret"})
        self.assertNotIn("Content-Type", kwargs["headers"])
        filename, field_value, content_type = kwargs["files"]["file"]
        self.assertEqual(filename, "reference_1.png")
        self.assertEqual(content_type, "image/png")
        self.assertTrue(field_value.startswith(b"\x89PNG"))

    async def test_original_jpeg_upload_uses_jpeg_multipart(self):
        expected_url = "https://cf-api.o1key.com/tmp/input/reference.jpg"
        response = _FakeResponse(200, {"url": expected_url})
        session = _FakeSession(response)
        image = _noise_image(16, 12)
        image._o1key_original_format = "JPEG"
        image._o1key_original_bytes = _jpeg_image_bytes(image)

        urls = await NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
            session=session,
            base_url="https://cf-api.o1key.com/",
            api_key="secret",
            images=[image],
        )

        self.assertEqual(urls, [expected_url])
        _, kwargs = session.calls[0]
        filename, field_value, content_type = kwargs["files"]["file"]
        self.assertEqual(filename, "reference_1.jpg")
        self.assertEqual(content_type, "image/jpeg")
        self.assertTrue(field_value.startswith(b"\xff\xd8"))

    async def test_concurrent_tasks_reuse_one_upload(self):
        image = Image.new("RGB", (2, 2), "blue")
        cache = {}
        calls = 0

        async def fake_upload(**_kwargs):
            nonlocal calls
            calls += 1
            await asyncio.sleep(0)
            return "https://example.invalid/reference.png"

        with patch.object(
            NANO_ASYNC,
            "_upload_nano_banana_temp_image",
            side_effect=fake_upload,
        ):
            first, second = await asyncio.gather(
                NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
                    session=object(),
                    base_url="https://example.invalid",
                    api_key="key",
                    images=[image],
                    upload_cache=cache,
                ),
                NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
                    session=object(),
                    base_url="https://example.invalid",
                    api_key="key",
                    images=[image],
                    upload_cache=cache,
                ),
            )

        self.assertEqual(calls, 1)
        self.assertEqual(first, second)

    async def test_all_reference_uploads_run_concurrently_and_keep_order(self):
        active = 0
        max_active = 0

        async def fake_upload(**kwargs):
            nonlocal active, max_active
            active += 1
            max_active = max(max_active, active)
            await asyncio.sleep(0.01)
            active -= 1
            return f"https://example.invalid/{kwargs['filename']}"

        images = [Image.new("RGB", (2, 2), color) for color in ("red", "green", "blue")]
        with patch.object(
            NANO_ASYNC,
            "_upload_nano_banana_temp_image",
            side_effect=fake_upload,
        ):
            urls = await NANO_ASYNC.upload_nano_banana_images_to_temp_urls(
                session=object(),
                base_url="https://example.invalid",
                api_key="key",
                images=images,
            )

        self.assertEqual(max_active, 3)
        self.assertEqual(urls, [
            "https://example.invalid/reference_1.png",
            "https://example.invalid/reference_2.png",
            "https://example.invalid/reference_3.png",
        ])

    async def test_generation_body_contains_inline_png_without_upload(self):
        image = Image.new("RGB", (2, 2), "green")
        result_image = Image.new("RGB", (2, 2), "white")

        with (
            patch.object(
                NANO_ASYNC,
                "upload_nano_banana_images_to_temp_urls",
                new=AsyncMock(),
            ) as upload,
            patch.object(
                NANO_ASYNC,
                "_submit_task",
                new=AsyncMock(return_value="task-1"),
            ) as submit,
            patch.object(
                NANO_ASYNC,
                "_poll_task",
                new=AsyncMock(return_value={"status": "success"}),
            ),
            patch.object(
                NANO_ASYNC,
                "_parse_task_images",
                new=AsyncMock(return_value=[result_image]),
            ),
        ):
            images, _ = await NANO_ASYNC.generate_nano_banana_async(
                session=object(),
                base_url="https://example.invalid",
                api_key="key",
                prompt="test",
                model="model",
                resolution="1K",
                aspect_ratio="1:1",
                images=[image],
            )

        self.assertEqual(images, [result_image])
        upload.assert_not_awaited()
        body = submit.await_args.args[3]
        self.assertEqual(len(body["images"]), 1)
        inline_data = body["images"][0]["inlineData"]
        self.assertEqual(inline_data["mimeType"], "image/png")
        self.assertTrue(base64.b64decode(inline_data["data"]).startswith(b"\x89PNG\r\n\x1a\n"))
        self.assertNotIn("data:image", inline_data["data"])
        sanitized = NANO_ASYNC._shorten_base64_for_log(body)
        self.assertEqual(
            sanitized["images"][0]["inlineData"]["data"],
            f"<base64 data, {len(inline_data['data'])} chars>",
        )

    def test_submit_body_declares_exact_jpeg_mime_type(self):
        image = _noise_image(16, 12)
        original = _jpeg_image_bytes(image)
        image._o1key_original_format = "JPEG"
        image._o1key_original_bytes = original

        body = NANO_ASYNC.build_nano_banana_submit_body(
            model="model",
            prompt="test",
            resolution="1K",
            aspect_ratio="1:1",
            images=[image],
        )

        inline_data = body["images"][0]["inlineData"]
        self.assertEqual(inline_data["mimeType"], "image/jpeg")
        self.assertEqual(base64.b64decode(inline_data["data"]), original)

    def test_submit_body_only_includes_google_search_when_enabled(self):
        disabled = NANO_ASYNC.build_nano_banana_submit_body(
            model="model",
            prompt="test",
            resolution="1K",
            aspect_ratio="1:1",
        )
        enabled = NANO_ASYNC.build_nano_banana_submit_body(
            model="model",
            prompt="test",
            resolution="1K",
            aspect_ratio="1:1",
            google_search=True,
        )

        self.assertNotIn("google_search", disabled)
        self.assertIs(enabled["google_search"], True)

    def test_submit_body_omits_size_for_smart_resolution(self):
        body = NANO_ASYNC.build_nano_banana_submit_body(
            model="model",
            prompt="test",
            resolution="智能",
            aspect_ratio="智能",
        )

        self.assertNotIn("size", body)

    def test_submit_body_rejects_reference_urls(self):
        with self.assertRaisesRegex(ValueError, "inlineData"):
            NANO_ASYNC.build_nano_banana_submit_body(
                model="model",
                prompt="test",
                resolution="1K",
                aspect_ratio="1:1",
                image_urls=["https://example.invalid/reference.png"],
            )

    async def test_no_resize_rejects_the_exact_json_body_above_the_local_limit(self):
        image = _noise_image(256, 128)
        with patch.object(NANO_ASYNC, "NANO_BANANA_REQUEST_BODY_LIMIT_BYTES", 20 * 1024):
            with self.assertRaisesRegex(ValueError, "当前设置为“不缩放”"):
                await NANO_ASYNC.prepare_nano_banana_inline_images(
                    [image],
                    model="model",
                    prompt="test",
                    resolution="1K",
                    aspect_ratio="2:1",
                    resize_mode="不缩放",
                )

    async def test_smart_resize_preserves_ratio_and_fits_the_exact_json_body(self):
        image = _noise_image(256, 128)
        with (
            patch.object(NANO_ASYNC, "NANO_BANANA_REQUEST_BODY_LIMIT_BYTES", 20 * 1024),
            patch.object(NANO_ASYNC, "_SMART_RESIZE_MIN_LONG_EDGE", 32),
        ):
            inline_images = await NANO_ASYNC.prepare_nano_banana_inline_images(
                [image],
                model="model",
                prompt="test",
                resolution="1K",
                aspect_ratio="2:1",
                resize_mode="智能缩放",
            )
            body = NANO_ASYNC.build_nano_banana_submit_body(
                model="model",
                prompt="test",
                resolution="1K",
                aspect_ratio="2:1",
                inline_images=inline_images,
            )
            body_size = NANO_ASYNC.validate_nano_banana_request_body(body)

        decoded = base64.b64decode(inline_images[0]["inlineData"]["data"])
        with Image.open(io.BytesIO(decoded)) as resized:
            self.assertLess(resized.width, image.width)
            self.assertAlmostEqual(resized.width / resized.height, 2.0, places=1)
        self.assertLessEqual(body_size, 20 * 1024)


if __name__ == "__main__":
    unittest.main(verbosity=2)