comfyui_o1key/tests/test_temp_media_uploads.py

import importlib.util
import json
import sys
import types
import unittest
from pathlib import Path
from unittest.mock import patch

import aiohttp
from PIL import Image


ROOT = Path(__file__).resolve().parents[1]


def _load_uploader():
    package = types.ModuleType("comfyui_o1key")
    package.__path__ = [str(ROOT)]
    utils_package = types.ModuleType("comfyui_o1key.utils")
    utils_package.__path__ = [str(ROOT / "utils")]
    sys.modules[package.__name__] = package
    sys.modules[utils_package.__name__] = utils_package

    config = types.ModuleType("comfyui_o1key.utils.config")
    config.get_api_key_or_raise = lambda *_args, **_kwargs: "secret"
    config.get_api_base_url = lambda: "https://api.o1key.cn"
    sys.modules[config.__name__] = config

    video_task = types.ModuleType("comfyui_o1key.utils.video_task")
    video_task.check_interrupt = lambda: None

    async def run_with_interrupt(coro):
        return await coro

    video_task.run_with_interrupt = run_with_interrupt
    sys.modules[video_task.__name__] = video_task

    spec = importlib.util.spec_from_file_location(
        "comfyui_o1key.utils.r2_uploader",
        ROOT / "utils" / "r2_uploader.py",
    )
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


UPLOADER = _load_uploader()


class _FakeResponse:
    def __init__(self, status, payload, headers=None):
        self.status = status
        self.payload = 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 json.dumps(self.payload)


class _FakeSession:
    def __init__(self, factory):
        self.factory = factory

    async def __aenter__(self):
        return self

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

    def post(self, url, **kwargs):
        self.factory.calls.append((url, kwargs))
        return self.factory.responses.pop(0)


class _SessionFactory:
    def __init__(self, responses):
        self.responses = list(responses)
        self.calls = []

    def __call__(self, **_kwargs):
        return _FakeSession(self)


class TempMediaUploadTests(unittest.IsolatedAsyncioTestCase):
    async def test_upload_uses_new_multipart_endpoint(self):
        expected_url = "https://cf-api.o1key.com/tmp/input/reference.png"
        factory = _SessionFactory([
            _FakeResponse(200, {
                "url": expected_url,
                "filename": "reference.png",
                "content_type": "image/png",
                "size": 12,
                "expires_at": 1_787_495_062,
            })
        ])

        with (
            patch.object(UPLOADER.aiohttp, "ClientSession", factory),
            patch.object(UPLOADER.aiohttp, "TCPConnector", return_value=object()),
            patch("builtins.print") as print_mock,
        ):
            result = await UPLOADER._upload_file(
                b"png-bytes",
                "reference.png",
                "image/png",
                base_url="https://cf-api.o1key.com/",
            )

        self.assertEqual(result, expected_url)
        rendered_log = "\n".join(
            " ".join(str(arg) for arg in call.args)
            for call in print_mock.call_args_list
        )
        self.assertNotIn(expected_url, rendered_log)
        self.assertEqual(len(factory.calls), 1)
        url, kwargs = factory.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"])
        self.assertIsInstance(kwargs["data"], aiohttp.FormData)
        field_options, field_headers, field_value = kwargs["data"]._fields[0]
        self.assertEqual(field_options["name"], "file")
        self.assertEqual(field_options["filename"], "reference.png")
        self.assertEqual(field_headers["Content-Type"], "image/png")
        self.assertEqual(field_value, b"png-bytes")

    async def test_non_retryable_upload_error_is_not_retried(self):
        factory = _SessionFactory([
            _FakeResponse(413, {"error": "attachment too large"})
        ])

        with (
            patch.object(UPLOADER.aiohttp, "ClientSession", factory),
            patch.object(UPLOADER.aiohttp, "TCPConnector", return_value=object()),
        ):
            with self.assertRaisesRegex(RuntimeError, "HTTP 413"):
                await UPLOADER._upload_file(
                    b"data",
                    "video.mp4",
                    "video/mp4",
                )

        self.assertEqual(len(factory.calls), 1)


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