comfyui_o1key/tests/test_http2_client.py

import importlib.util
import sys
import unittest
from pathlib import Path
from unittest.mock import MagicMock, patch


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


def _load_module():
    spec = importlib.util.spec_from_file_location(
        "o1key_http2_client_test_module",
        ROOT / "utils" / "http2_client.py",
    )
    module = importlib.util.module_from_spec(spec)
    sys.modules[spec.name] = module
    spec.loader.exec_module(module)
    return module


def _load_module_without_httpx():
    module_name = "o1key_http2_client_without_httpx_test_module"
    spec = importlib.util.spec_from_file_location(
        module_name,
        ROOT / "utils" / "http2_client.py",
    )
    module = importlib.util.module_from_spec(spec)
    with patch.dict(sys.modules, {"httpx": None}):
        sys.modules[module_name] = module
        spec.loader.exec_module(module)
    return module


HTTP2_CLIENT = _load_module()


class _ChunkContent:
    def __init__(self, chunks, error=None):
        self.chunks = chunks
        self.error = error

    async def iter_chunked(self, _chunk_size):
        for chunk in self.chunks:
            yield chunk
        if self.error is not None:
            raise self.error


class _TraceResponse:
    def __init__(self, chunks, headers=None, error=None):
        self.status = 200
        self.http_version = "HTTP/2"
        self.headers = headers or {}
        self.content = _ChunkContent(chunks, error=error)


class Http2ClientTests(unittest.TestCase):
    def test_enables_http2_when_runtime_support_is_present(self):
        fake_client = MagicMock()
        with (
            patch.object(HTTP2_CLIENT, "http2_runtime_available", return_value=True),
            patch.object(HTTP2_CLIENT.httpx, "AsyncClient", return_value=fake_client) as constructor,
        ):
            client = HTTP2_CLIENT.O1keyAsyncHttpClient(http2=True)

        self.assertTrue(client.http2_enabled)
        self.assertTrue(constructor.call_args.kwargs["http2"])
        self.assertTrue(constructor.call_args.kwargs["verify"])

    def test_falls_back_to_http11_when_h2_runtime_is_missing(self):
        fake_client = MagicMock()
        with (
            patch.object(HTTP2_CLIENT, "http2_runtime_available", return_value=False),
            patch.object(HTTP2_CLIENT.httpx, "AsyncClient", return_value=fake_client) as constructor,
        ):
            client = HTTP2_CLIENT.O1keyAsyncHttpClient(http2=True)

        self.assertFalse(client.http2_enabled)
        self.assertFalse(constructor.call_args.kwargs["http2"])

    def test_task_id_validation_uses_only_explicit_task_fields(self):
        payload = {"id": "result-image-id", "data": {"taskId": "task-7"}}
        self.assertEqual(HTTP2_CLIENT.response_task_id(payload), "task-7")
        self.assertEqual(
            HTTP2_CLIENT.validate_response_task_id(payload, "task-7"),
            "task-7",
        )
        with self.assertRaisesRegex(
            HTTP2_CLIENT.ResponseTaskIdMismatchError,
            "requested_task_id=task-8.*response_task_id=task-7",
        ):
            HTTP2_CLIENT.validate_response_task_id(payload, "task-8")


class ResponseBodyDiagnosticsTests(unittest.IsolatedAsyncioTestCase):
    async def test_exact_content_length_is_reported_as_match(self):
        response = _TraceResponse(
            [b'{"ok":', b'true}'],
            headers={"Content-Length": "11"},
        )

        body, diagnostics = await HTTP2_CLIENT.read_response_body_with_diagnostics(response)

        self.assertEqual(body, b'{"ok":true}')
        self.assertEqual(diagnostics["declared_bytes"], 11)
        self.assertEqual(diagnostics["received_bytes"], 11)
        self.assertEqual(diagnostics["length_check"], "match")

    async def test_short_content_length_raises_with_received_byte_count(self):
        response = _TraceResponse(
            [b"1234"],
            headers={"Content-Length": "10"},
        )

        with self.assertRaisesRegex(
            HTTP2_CLIENT.ResponseBodyIntegrityError,
            r"Content-Length=10.*received=4B.*length_check=mismatch",
        ):
            await HTTP2_CLIENT.read_response_body_with_diagnostics(response)

    async def test_stream_failure_keeps_partial_received_byte_count(self):
        response = _TraceResponse(
            [b"1234"],
            headers={"Content-Length": "10"},
            error=OSError("connection closed"),
        )

        with self.assertRaisesRegex(
            HTTP2_CLIENT.ResponseBodyIntegrityError,
            r"读取提前中断.*Content-Length=10.*received=4B.*connection closed",
        ):
            await HTTP2_CLIENT.read_response_body_with_diagnostics(response)

    async def test_compressed_response_does_not_compare_decoded_size(self):
        response = _TraceResponse(
            [b"decoded body"],
            headers={"Content-Length": "5", "Content-Encoding": "gzip"},
        )

        _body, diagnostics = await HTTP2_CLIENT.read_response_body_with_diagnostics(response)

        self.assertEqual(diagnostics["length_check"], "skipped-compressed")


class AiohttpFallbackTests(unittest.IsolatedAsyncioTestCase):
    async def test_missing_httpx_uses_working_aiohttp_session(self):
        module = _load_module_without_httpx()

        self.assertFalse(module.HTTPX_AVAILABLE)
        self.assertFalse(module.http2_runtime_available())
        self.assertIsInstance(
            module.create_timeout(
                120.0,
                connect=30.0,
                read=60.0,
                write=30.0,
                pool=30.0,
            ),
            module.aiohttp.ClientTimeout,
        )

        client = module.O1keyAsyncHttpClient(http2=True)
        self.assertEqual(client.backend, "aiohttp")
        self.assertFalse(client.http2_enabled)
        async with client as active_client:
            self.assertIsInstance(active_client._client, module.aiohttp.ClientSession)

    async def test_missing_httpx_converts_files_to_aiohttp_multipart(self):
        module = _load_module_without_httpx()
        client = module.O1keyAsyncHttpClient(http2=True)
        fake_session = MagicMock()
        sentinel_context = object()
        fake_session.post.return_value = sentinel_context
        client._client = fake_session

        result = client.post(
            "https://example.invalid/upload",
            headers={"Authorization": "Bearer test"},
            files={"file": ("reference.jpg", b"jpeg", "image/jpeg")},
            timeout=module.create_timeout(
                120.0,
                connect=30.0,
                read=60.0,
                write=30.0,
                pool=30.0,
            ),
        )

        self.assertIs(result, sentinel_context)
        payload = fake_session.post.call_args.kwargs["data"]
        self.assertIsInstance(payload, module.aiohttp.FormData)
        self.assertEqual(len(payload._fields), 1)


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