| import http.client |
| import base64 |
| import json |
| import threading |
| import unittest |
| from unittest import mock |
| from urllib.parse import parse_qs |
|
|
| from gemini_web2api.config import CONFIG, DEFAULT_CONFIG |
| from gemini_web2api.gemini import _build_payload |
| from gemini_web2api.server import GeminiHandler, ThreadedServer |
| from gemini_web2api.tools import google_contents_to_prompt, messages_to_prompt |
|
|
|
|
| def _decode_payload(payload): |
| outer = json.loads(parse_qs(payload)["f.req"][0]) |
| return json.loads(outer[1]) |
|
|
|
|
| def _decode_sse(body): |
| events = [] |
| for block in body.strip().split("\n\n"): |
| lines = block.splitlines() |
| event_type = next( |
| (line[len("event: "):] for line in lines if line.startswith("event: ")), |
| None, |
| ) |
| data = next( |
| (line[len("data: "):] for line in lines if line.startswith("data: ")), |
| None, |
| ) |
| if event_type and data: |
| events.append((event_type, json.loads(data))) |
| return events |
|
|
|
|
| class PayloadPersistenceTests(unittest.TestCase): |
| def setUp(self): |
| self.original_config = dict(CONFIG) |
|
|
| def tearDown(self): |
| CONFIG.clear() |
| CONFIG.update(self.original_config) |
|
|
| def test_temporary_chats_default_to_disabled(self): |
| self.assertIs(DEFAULT_CONFIG["temporary_chats"], False) |
|
|
| def test_persistent_chat_payload(self): |
| CONFIG["temporary_chats"] = False |
|
|
| inner = _decode_payload(_build_payload("hello", 1, 4)) |
|
|
| self.assertEqual(inner[41], [2]) |
| self.assertIsNone(inner[45]) |
|
|
| def test_temporary_chat_payload(self): |
| CONFIG["temporary_chats"] = True |
|
|
| inner = _decode_payload(_build_payload("hello", 1, 4)) |
|
|
| self.assertEqual(inner[41], [1]) |
| self.assertEqual(inner[45], 1) |
|
|
| def test_payload_includes_uploaded_image_refs(self): |
| inner = _decode_payload(_build_payload("describe", 1, 4, ["/uploaded/image-ref"])) |
|
|
| self.assertEqual(inner[0][0], "describe") |
| self.assertEqual(inner[0][3], [[None, None, "/uploaded/image-ref"]]) |
|
|
|
|
| class MessageParsingTests(unittest.TestCase): |
| def test_messages_to_prompt_extracts_openai_image_url_data_url(self): |
| image_data = base64.b64encode(b"fake png").decode() |
|
|
| prompt, images = messages_to_prompt([{ |
| "role": "user", |
| "content": [ |
| {"type": "text", "text": "Describe"}, |
| {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_data}"}}, |
| ], |
| }]) |
|
|
| self.assertEqual(prompt, "Describe [Image attached]") |
| self.assertEqual(images, [(b"fake png", "image/png")]) |
|
|
| def test_messages_to_prompt_extracts_responses_input_image_url(self): |
| prompt, images = messages_to_prompt([{ |
| "role": "user", |
| "content": [ |
| {"type": "input_text", "text": "Describe"}, |
| {"type": "input_image", "image_url": "https://example.com/image.png"}, |
| ], |
| }]) |
|
|
| self.assertEqual(prompt, "Describe [Image attached]") |
| self.assertEqual(images, [("https://example.com/image.png", "image/png")]) |
|
|
| def test_messages_to_prompt_ignores_malformed_image_data_url(self): |
| prompt, images = messages_to_prompt([{ |
| "role": "user", |
| "content": [ |
| {"type": "text", "text": "Describe"}, |
| {"type": "image_url", "image_url": {"url": "data:image/png;base64,%%%"}}, |
| ], |
| }]) |
|
|
| self.assertEqual(prompt, "Describe") |
| self.assertEqual(images, []) |
|
|
| def test_google_contents_to_prompt_extracts_inline_image_data(self): |
| image_data = base64.b64encode(b"fake png").decode() |
|
|
| prompt, images = google_contents_to_prompt({ |
| "contents": [{ |
| "role": "user", |
| "parts": [ |
| {"text": "Describe"}, |
| {"inlineData": {"mimeType": "image/png", "data": image_data}}, |
| ], |
| }], |
| }) |
|
|
| self.assertEqual(prompt, "Describe\n[Image attached]") |
| self.assertEqual(images, [(b"fake png", "image/png")]) |
|
|
| def test_google_contents_to_prompt_ignores_malformed_inline_image_data(self): |
| prompt, images = google_contents_to_prompt({ |
| "contents": [{ |
| "role": "user", |
| "parts": [ |
| {"text": "Describe"}, |
| {"inlineData": {"mimeType": "image/png", "data": "%%%"}}, |
| ], |
| }], |
| }) |
|
|
| self.assertEqual(prompt, "Describe") |
| self.assertEqual(images, []) |
|
|
|
|
| class StreamingEndpointTests(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| cls.server = ThreadedServer(("127.0.0.1", 0), GeminiHandler) |
| cls.thread = threading.Thread(target=cls.server.serve_forever, daemon=True) |
| cls.thread.start() |
| cls.port = cls.server.server_address[1] |
|
|
| @classmethod |
| def tearDownClass(cls): |
| cls.server.shutdown() |
| cls.server.server_close() |
| cls.thread.join(timeout=5) |
|
|
| def setUp(self): |
| self.original_config = dict(CONFIG) |
| CONFIG["api_keys"] = [] |
| CONFIG["log_requests"] = False |
|
|
| def tearDown(self): |
| CONFIG.clear() |
| CONFIG.update(self.original_config) |
|
|
| def post_json(self, path, payload): |
| connection = http.client.HTTPConnection("127.0.0.1", self.port, timeout=5) |
| connection.request( |
| "POST", |
| path, |
| body=json.dumps(payload), |
| headers={"Content-Type": "application/json"}, |
| ) |
| response = connection.getresponse() |
| body = response.read().decode() |
| headers = dict(response.getheaders()) |
| connection.close() |
| return response.status, headers, body |
|
|
| def post_chunked_json(self, path, payload): |
| connection = http.client.HTTPConnection("127.0.0.1", self.port, timeout=5) |
| connection.request( |
| "POST", |
| path, |
| body=json.dumps(payload).encode(), |
| headers={"Content-Type": "application/json"}, |
| encode_chunked=True, |
| ) |
| response = connection.getresponse() |
| body = response.read().decode() |
| headers = dict(response.getheaders()) |
| connection.close() |
| return response.status, headers, body |
|
|
| @mock.patch("gemini_web2api.server.generate_stream") |
| def test_chat_stream_starts_with_assistant_role(self, generate_stream): |
| generate_stream.return_value = iter(["hel", "lo"]) |
|
|
| status, headers, body = self.post_json( |
| "/v1/chat/completions", |
| { |
| "model": "gemini-3.6-flash", |
| "messages": [{"role": "user", "content": "hello"}], |
| "stream": True, |
| }, |
| ) |
|
|
| self.assertEqual(status, 200) |
| self.assertEqual(headers["Content-Type"], "text/event-stream") |
| chunks = [ |
| json.loads(line[len("data: "):]) |
| for line in body.splitlines() |
| if line.startswith("data: {") |
| ] |
| self.assertEqual(chunks[0]["choices"][0]["delta"], {"role": "assistant"}) |
| self.assertEqual(chunks[1]["choices"][0]["delta"], {"content": "hel"}) |
| self.assertEqual(chunks[2]["choices"][0]["delta"], {"content": "lo"}) |
| self.assertTrue(body.endswith("data: [DONE]\n\n")) |
|
|
| @mock.patch("gemini_web2api.server.generate", return_value="chunked ok") |
| def test_chat_accepts_chunked_body(self, _generate): |
| status, _, body = self.post_chunked_json( |
| "/v1/chat/completions", |
| { |
| "model": "gemini-3.6-flash", |
| "messages": [{"role": "user", "content": "hello"}], |
| }, |
| ) |
|
|
| self.assertEqual(status, 200) |
| self.assertEqual(json.loads(body)["choices"][0]["message"]["content"], "chunked ok") |
|
|
| @mock.patch("gemini_web2api.server.upload_image", return_value="/uploaded/image-ref") |
| @mock.patch("gemini_web2api.server.generate", return_value="looks good") |
| def test_chat_accepts_openai_image_url_data_url(self, generate, upload_image): |
| image_data = base64.b64encode(b"fake png").decode() |
|
|
| status, _, body = self.post_json( |
| "/v1/chat/completions", |
| { |
| "model": "gemini-3.6-flash", |
| "messages": [{ |
| "role": "user", |
| "content": [ |
| {"type": "text", "text": "Describe this image"}, |
| { |
| "type": "image_url", |
| "image_url": { |
| "url": f"data:image/png;base64,{image_data}" |
| }, |
| }, |
| ], |
| }], |
| }, |
| ) |
|
|
| self.assertEqual(status, 200) |
| upload_image.assert_called_once_with(b"fake png", "image.png", "image/png") |
| self.assertEqual(generate.call_args.args[3], ["/uploaded/image-ref"]) |
| self.assertIn("[Image attached]", generate.call_args.args[0]) |
| self.assertEqual(json.loads(body)["choices"][0]["message"]["content"], "looks good") |
|
|
| @mock.patch("gemini_web2api.server.fetch_image_bytes", return_value=b"\xff\xd8\xffremote jpeg") |
| @mock.patch("gemini_web2api.server.upload_image", return_value="/uploaded/remote-ref") |
| @mock.patch("gemini_web2api.server.generate", return_value="remote ok") |
| def test_responses_accepts_input_image_url(self, generate, upload_image, fetch_image_bytes): |
| status, _, _ = self.post_json( |
| "/v1/responses", |
| { |
| "model": "gemini-3.6-flash", |
| "input": [{ |
| "role": "user", |
| "content": [ |
| {"type": "input_text", "text": "What is shown?"}, |
| { |
| "type": "input_image", |
| "image_url": "https://example.com/image.jpg", |
| }, |
| ], |
| }], |
| }, |
| ) |
|
|
| self.assertEqual(status, 200) |
| fetch_image_bytes.assert_called_once_with("https://example.com/image.jpg") |
| upload_image.assert_called_once_with(b"\xff\xd8\xffremote jpeg", "image.png", "image/jpeg") |
| self.assertEqual(generate.call_args.args[3], ["/uploaded/remote-ref"]) |
| self.assertIn("[Image attached]", generate.call_args.args[0]) |
|
|
| @mock.patch("gemini_web2api.server.upload_image", return_value="/uploaded/image-ref") |
| @mock.patch("gemini_web2api.server.generate", return_value="top-level image ok") |
| def test_responses_accepts_top_level_input_image(self, generate, upload_image): |
| image_data = base64.b64encode(b"fake png").decode() |
|
|
| status, _, _ = self.post_json( |
| "/v1/responses", |
| { |
| "model": "gemini-3.6-flash", |
| "input": [ |
| {"type": "input_text", "text": "What is shown?"}, |
| { |
| "type": "input_image", |
| "image_url": f"data:image/png;base64,{image_data}", |
| }, |
| ], |
| }, |
| ) |
|
|
| self.assertEqual(status, 200) |
| upload_image.assert_called_once_with(b"fake png", "image.png", "image/png") |
| self.assertEqual(generate.call_args.args[3], ["/uploaded/image-ref"]) |
| self.assertIn("What is shown?", generate.call_args.args[0]) |
| self.assertIn("[Image attached]", generate.call_args.args[0]) |
|
|
| @mock.patch("gemini_web2api.server.upload_image", side_effect=RuntimeError("upload denied")) |
| def test_google_image_upload_failure_returns_502(self, _upload_image): |
| image_data = base64.b64encode(b"fake png").decode() |
|
|
| status, _, body = self.post_json( |
| "/v1beta/models/gemini-3.6-flash:generateContent", |
| { |
| "contents": [{ |
| "role": "user", |
| "parts": [{ |
| "inlineData": { |
| "mimeType": "image/png", |
| "data": image_data, |
| }, |
| }], |
| }], |
| }, |
| ) |
|
|
| self.assertEqual(status, 502) |
| self.assertIn("image upload failed: upload denied", json.loads(body)["error"]["message"]) |
|
|
| @mock.patch("gemini_web2api.server.generate_stream", return_value=iter(["streamed"])) |
| def test_google_stream_generate_content_uses_sse(self, _generate_stream): |
| status, headers, body = self.post_json( |
| "/v1beta/models/gemini-3.6-flash:streamGenerateContent", |
| { |
| "contents": [{ |
| "role": "user", |
| "parts": [{"text": "Stream this"}], |
| }], |
| }, |
| ) |
|
|
| self.assertEqual(status, 200) |
| self.assertEqual(headers["Content-Type"], "text/event-stream") |
| self.assertIn('"text": "streamed"', body) |
|
|
| @mock.patch("gemini_web2api.server.generate", return_value="hello") |
| def test_responses_text_stream_has_complete_event_sequence(self, _generate): |
| status, headers, body = self.post_json( |
| "/v1/responses", |
| { |
| "model": "gemini-3.6-flash", |
| "input": "hello", |
| "stream": True, |
| }, |
| ) |
|
|
| self.assertEqual(status, 200) |
| self.assertEqual(headers["Content-Type"], "text/event-stream") |
| events = _decode_sse(body) |
| self.assertEqual( |
| [event_type for event_type, _ in events], |
| [ |
| "response.created", |
| "response.in_progress", |
| "response.output_item.added", |
| "response.content_part.added", |
| "response.output_text.delta", |
| "response.output_text.done", |
| "response.content_part.done", |
| "response.output_item.done", |
| "response.completed", |
| ], |
| ) |
| self.assertEqual( |
| [event["sequence_number"] for _, event in events], |
| list(range(1, len(events) + 1)), |
| ) |
| self.assertEqual(events[4][1]["delta"], "hello") |
| self.assertEqual(events[-1][1]["response"]["status"], "completed") |
| self.assertEqual(events[-1][1]["response"]["output"][0]["content"][0]["text"], "hello") |
|
|
| @mock.patch("gemini_web2api.server.parse_tool_calls") |
| @mock.patch("gemini_web2api.server.generate", return_value="tool output") |
| def test_responses_function_call_stream_has_complete_event_sequence( |
| self, _generate, parse_tool_calls |
| ): |
| parse_tool_calls.return_value = ( |
| "", |
| [ |
| { |
| "id": "call_test", |
| "type": "function", |
| "function": {"name": "get_weather", "arguments": '{"city":"Shanghai"}'}, |
| } |
| ], |
| ) |
|
|
| status, _, body = self.post_json( |
| "/v1/responses", |
| { |
| "model": "gemini-3.6-flash", |
| "input": "weather", |
| "tools": [ |
| { |
| "type": "function", |
| "name": "get_weather", |
| "description": "Get weather", |
| "parameters": {"type": "object"}, |
| } |
| ], |
| "stream": True, |
| }, |
| ) |
|
|
| self.assertEqual(status, 200) |
| events = _decode_sse(body) |
| self.assertEqual( |
| [event_type for event_type, _ in events], |
| [ |
| "response.created", |
| "response.in_progress", |
| "response.output_item.added", |
| "response.function_call_arguments.delta", |
| "response.function_call_arguments.done", |
| "response.output_item.done", |
| "response.completed", |
| ], |
| ) |
| self.assertEqual( |
| [event["sequence_number"] for _, event in events], |
| list(range(1, len(events) + 1)), |
| ) |
| self.assertEqual(events[2][1]["output_index"], 0) |
| self.assertEqual(events[3][1]["delta"], '{"city":"Shanghai"}') |
| self.assertEqual(events[4][1]["arguments"], '{"city":"Shanghai"}') |
| self.assertEqual(events[-1][1]["response"]["output"][0]["name"], "get_weather") |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
