From f3a65bd26b1828e3fc1c7271115d6de5a06b4047 Mon Sep 17 00:00:00 2001 From: CrystalVibe28 <82704512+CrystalVibe28@users.noreply.github.com> Date: Thu, 14 May 2026 22:17:07 +0800 Subject: [PATCH] test(chatgpt): cover image generation transformation --- .../test_chatgpt_image_generation.py | 397 +++++++++++++++++- 1 file changed, 391 insertions(+), 6 deletions(-) diff --git a/tests/image_gen_tests/test_chatgpt_image_generation.py b/tests/image_gen_tests/test_chatgpt_image_generation.py index 2926d495bf5..495a62a3ea8 100644 --- a/tests/image_gen_tests/test_chatgpt_image_generation.py +++ b/tests/image_gen_tests/test_chatgpt_image_generation.py @@ -1,9 +1,14 @@ +from typing import Any, cast + import httpx import pytest +from litellm.exceptions import AuthenticationError +from litellm.llms.chatgpt.common_utils import GetAccessTokenError from litellm.llms.chatgpt.image_generation.transformation import ( ChatGPTImageGenerationConfig, ) +from litellm.llms.openai.common_utils import OpenAIError from litellm.types.utils import LlmProviders from litellm.types.utils import ImageResponse from litellm.utils import ProviderConfigManager @@ -14,6 +19,10 @@ class MockLogging: pass +def mock_logging() -> Any: + return MockLogging() + + def test_chatgpt_image_generation_transforms_request(monkeypatch, tmp_path): monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) config = ChatGPTImageGenerationConfig() @@ -157,6 +166,147 @@ def test_chatgpt_image_generation_config_registered(monkeypatch, tmp_path): assert isinstance(config, ChatGPTImageGenerationConfig) +def test_chatgpt_image_generation_maps_supported_openai_params(monkeypatch, tmp_path): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + optional_params = {"quality": "high"} + result = config.map_openai_params( + non_default_params={ + "quality": "low", + "size": "1024x1024", + "unsupported": "drop-me", + }, + optional_params=optional_params, + model="gpt-image-2", + drop_params=True, + ) + + assert result is optional_params + assert result == {"quality": "high", "size": "1024x1024"} + + +def test_chatgpt_image_generation_rejects_unsupported_openai_param( + monkeypatch, tmp_path +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + with pytest.raises(ValueError, match="Parameter unsupported is not supported"): + config.map_openai_params( + non_default_params={"unsupported": "keep-me"}, + optional_params={}, + model="gpt-image-2", + drop_params=False, + ) + + +def test_chatgpt_image_generation_validates_environment(monkeypatch, tmp_path): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + class FakeAuthenticator: + def get_access_token(self): + return "access-token" + + def get_account_id(self): + return "account-id" + + config.authenticator = cast(Any, FakeAuthenticator()) + + headers = config.validate_environment( + headers={"content-type": "application/custom", "x-extra": "1"}, + model="gpt-image-2", + messages=[], + optional_params={}, + litellm_params={"session_id": "session-123"}, + ) + + assert headers["Authorization"] == "Bearer access-token" + assert headers["ChatGPT-Account-Id"] == "account-id" + assert headers["session_id"] == "session-123" + assert headers["content-type"] == "application/custom" + assert headers["x-extra"] == "1" + + +def test_chatgpt_image_generation_validate_environment_auth_error( + monkeypatch, tmp_path +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + class FakeAuthenticator: + def get_access_token(self): + raise GetAccessTokenError(status_code=401, message="token expired") + + config.authenticator = cast(Any, FakeAuthenticator()) + + with pytest.raises(AuthenticationError, match="token expired"): + config.validate_environment( + headers={}, + model="gpt-image-2", + messages=[], + optional_params={}, + litellm_params={}, + ) + + +@pytest.mark.parametrize( + "api_base, expected", + [ + ( + "https://chatgpt.com/backend-api", + "https://chatgpt.com/backend-api/codex/responses", + ), + ( + "https://chatgpt.com/backend-api/responses", + "https://chatgpt.com/backend-api/codex/responses", + ), + ("https://example.test/custom/", "https://example.test/custom/responses"), + ], +) +def test_chatgpt_image_generation_get_complete_url_canonicalizes_api_base( + monkeypatch, tmp_path, api_base, expected +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + assert ( + config.get_complete_url( + api_base=api_base, + api_key=None, + model="gpt-image-2", + optional_params={}, + litellm_params={}, + ) + == expected + ) + + +def test_chatgpt_image_generation_get_complete_url_uses_authenticator_api_base( + monkeypatch, tmp_path +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + class FakeAuthenticator: + def get_api_base(self): + return "https://example.test/backend-api" + + config.authenticator = cast(Any, FakeAuthenticator()) + + assert ( + config.get_complete_url( + api_base=None, + api_key=None, + model="gpt-image-2", + optional_params={}, + litellm_params={}, + ) + == "https://example.test/backend-api/codex/responses" + ) + + def test_chatgpt_image_generation_extracts_b64_from_sse_completed_response( monkeypatch, tmp_path ): @@ -176,20 +326,97 @@ def test_chatgpt_image_generation_extracts_b64_from_sse_completed_response( model="gpt-image-2", raw_response=raw_response, model_response=ImageResponse(), - logging_obj=MockLogging(), + logging_obj=mock_logging(), request_data={"input": "draw a cat"}, optional_params={"size": "1024x1024", "quality": "high"}, litellm_params={}, encoding=None, ) + assert response.data is not None assert response.data[0].b64_json == "b64-image-data" assert response.size == "1024x1024" assert response.quality == "high" + assert response.output_format is None assert response.usage is None + assert response._hidden_params is not None assert response._hidden_params["model"] == "gpt-image-2" +def test_chatgpt_image_generation_uses_optional_responses_model_and_env( + monkeypatch, tmp_path +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setenv("CHATGPT_IMAGE_RESPONSES_MODEL", "gpt-env") + config = ChatGPTImageGenerationConfig() + optional_params = {"chatgpt_responses_model": "gpt-override"} + + request = config.transform_image_generation_request( + model="gpt-image-2", + prompt="draw a cat", + optional_params=optional_params, + litellm_params={"chatgpt_responses_model": "gpt-litellm-param"}, + headers={}, + ) + + assert request["model"] == "gpt-override" + assert "chatgpt_responses_model" not in optional_params + + env_request = config.transform_image_generation_request( + model="gpt-image-2", + prompt="draw a cat", + optional_params={}, + litellm_params={}, + headers={}, + ) + assert env_request["model"] == "gpt-env" + + +@pytest.mark.parametrize( + "model, optional_params, error", + [ + ("dall-e-3", {}, "requires a GPT Image model"), + ("gpt-image-1.5", {"size": "auto"}, None), + ("gpt-image-2", {"size": "auto"}, None), + ("gpt-image-2", {"size": "bad-size"}, "size must be auto or WIDTHxHEIGHT"), + ( + "gpt-image-1.5", + {"background": "transparent", "output_format": "jpeg"}, + "transparent background requires output_format png or webp", + ), + ( + "gpt-image-1.5", + {"background": "not-real"}, + "background must be one of transparent, opaque, or auto", + ), + ], +) +def test_chatgpt_image_generation_validates_additional_param_paths( + monkeypatch, tmp_path, model, optional_params, error +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + if error is None: + config.transform_image_generation_request( + model=model, + prompt="draw a cat", + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + return + + with pytest.raises(ValueError, match=error): + config.transform_image_generation_request( + model=model, + prompt="draw a cat", + optional_params=optional_params, + litellm_params={}, + headers={}, + ) + + def test_chatgpt_image_generation_extracts_b64_from_deep_nested_payload( monkeypatch, tmp_path ): @@ -205,6 +432,147 @@ def test_chatgpt_image_generation_extracts_b64_from_deep_nested_payload( assert partial_images == [] +def test_chatgpt_image_generation_extracts_json_response(monkeypatch, tmp_path): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + raw_response = httpx.Response( + status_code=200, + json={ + "output": [ + { + "type": "image_generation", + "image": ["b64-image-data", "b64-image-data", 123], + } + ], + "tool_usage": { + "image_gen": { + "input_tokens": 11, + "input_tokens_details": {"image_tokens": 1, "text_tokens": 10}, + "output_tokens": 22, + } + }, + }, + ) + + response = config.transform_image_generation_response( + model="gpt-image-2", + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=mock_logging(), + request_data={}, + optional_params={"response_format": "b64_json"}, + litellm_params={}, + encoding=None, + ) + + assert response.data is not None + assert [item.b64_json for item in response.data] == ["b64-image-data"] + assert response.output_format == "b64_json" + assert response.usage is not None + assert response.usage.input_tokens == 11 + assert response.usage.input_tokens_details.image_tokens == 1 + assert response.usage.input_tokens_details.text_tokens == 10 + assert response.usage.output_tokens == 22 + assert response.usage.total_tokens == 33 + + +def test_chatgpt_image_generation_raises_when_no_image_data(monkeypatch, tmp_path): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + raw_response = httpx.Response(status_code=200, json={"output": []}) + + with pytest.raises(OpenAIError, match="No image data found"): + config.transform_image_generation_response( + model="gpt-image-2", + raw_response=raw_response, + model_response=ImageResponse(), + logging_obj=mock_logging(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_chatgpt_image_generation_raises_provider_error_event(monkeypatch, tmp_path): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + with pytest.raises(OpenAIError, match="image blocked"): + config._extract_images_from_payload( + { + "type": "response.failed", + "response": {"error": {"message": "image blocked"}}, + } + ) + + +def test_chatgpt_image_generation_handles_invalid_json_payloads(monkeypatch, tmp_path): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + raw_response = httpx.Response( + status_code=200, + headers={"content-type": "application/json"}, + text="{not-json", + ) + + assert config._extract_image_payloads(raw_response) == [] + assert config._get_parsed_payloads(raw_response) == [{}] + + +def test_chatgpt_image_generation_ignores_non_dict_json_payload(monkeypatch, tmp_path): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + raw_response = httpx.Response(status_code=200, json=[]) + + assert config._extract_image_payloads(raw_response) == [] + assert config._get_parsed_payloads(raw_response) == [] + + +def test_chatgpt_image_generation_extracts_from_cyclic_nested_payload( + monkeypatch, tmp_path +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + payload = { + "type": "image_generation_call", + "result": "b64-result", + "b64_json": "b64-json", + "image": ["b64-image", 123], + } + payload["self"] = payload + + assert config._extract_images_from_nested_value(payload) == [ + "b64-result", + "b64-json", + "b64-image", + ] + + cyclic_list = [] + cyclic_list.append(cyclic_list) + assert config._extract_images_from_nested_value(cyclic_list) == [] + + +def test_chatgpt_image_generation_usage_helpers_ignore_invalid_payloads( + monkeypatch, tmp_path +): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + assert config._get_image_generation_usage("not-a-dict") is None + assert config._get_image_generation_usage({"tool_usage": []}) is None + assert config._get_image_generation_usage({"tool_usage": {"image_gen": []}}) is None + assert ( + config._get_image_generation_usage( + {"tool_usage": {"image_gen": {"input_tokens": 1}}} + ) + is None + ) + assert config._is_zero_image_usage( + {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0} + ) + + def test_chatgpt_image_generation_extracts_tool_usage_from_completed_response( monkeypatch, tmp_path ): @@ -230,7 +598,7 @@ def test_chatgpt_image_generation_extracts_tool_usage_from_completed_response( model="gpt-image-2", raw_response=raw_response, model_response=ImageResponse(), - logging_obj=MockLogging(), + logging_obj=mock_logging(), request_data={"input": "draw a cat"}, optional_params={}, litellm_params={}, @@ -286,7 +654,7 @@ def test_chatgpt_image_generation_prefers_completed_tool_usage(monkeypatch, tmp_ model="gpt-image-2", raw_response=raw_response, model_response=ImageResponse(), - logging_obj=MockLogging(), + logging_obj=mock_logging(), request_data={"input": "draw a cat"}, optional_params={}, litellm_params={}, @@ -332,13 +700,14 @@ def test_chatgpt_image_generation_extracts_usage_with_partial_image_payload( model="gpt-image-2", raw_response=raw_response, model_response=ImageResponse(), - logging_obj=MockLogging(), + logging_obj=mock_logging(), request_data={"input": "draw a cat"}, optional_params={}, litellm_params={}, encoding=None, ) + assert response.data is not None assert response.data[0].b64_json == "partial-image-data" assert response.usage is not None assert response.usage.input_tokens == 105 @@ -368,7 +737,7 @@ def test_chatgpt_image_generation_extracts_top_level_tool_usage(monkeypatch, tmp model="gpt-image-2", raw_response=raw_response, model_response=ImageResponse(), - logging_obj=MockLogging(), + logging_obj=mock_logging(), request_data={"input": "draw a cat"}, optional_params={}, litellm_params={}, @@ -402,11 +771,27 @@ def test_chatgpt_image_generation_extracts_b64_from_streaming_completed_event( model="gpt-image-2", raw_response=raw_response, model_response=ImageResponse(), - logging_obj=MockLogging(), + logging_obj=mock_logging(), request_data={"input": "draw a cat"}, optional_params={}, litellm_params={}, encoding=None, ) + assert response.data is not None assert [item.b64_json for item in response.data] == ["final-image"] + + +def test_chatgpt_image_generation_get_error_class(monkeypatch, tmp_path): + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + config = ChatGPTImageGenerationConfig() + + error = config.get_error_class( + error_message="bad request", + status_code=400, + headers={"x-request-id": "req-123"}, + ) + + assert isinstance(error, OpenAIError) + assert error.status_code == 400 + assert error.message == "bad request"