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)