From 2b13c9aba26fd37a6a246a9ce7900a8cf605a588 Mon Sep 17 00:00:00 2001
From: Sameer Kankute <sameer@berri.ai>
Date: Wed, 14 Jan 2026 10:17:01 +0530
Subject: [PATCH] Add tests for openrouter

---
 .../image_generation/transformation.py        |   3 +-
 ...est_openrouter_image_gen_transformation.py | 573 ++++++++++++++++++
 2 files changed, 575 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py

diff --git a/litellm/llms/openrouter/image_generation/transformation.py b/litellm/llms/openrouter/image_generation/transformation.py
index 737e857b03..a6790870f0 100644
--- a/litellm/llms/openrouter/image_generation/transformation.py
+++ b/litellm/llms/openrouter/image_generation/transformation.py
@@ -32,6 +32,7 @@ from typing import TYPE_CHECKING, Any, List, Optional, Union
 import httpx
 
 import litellm
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
 from litellm.llms.base_llm.image_generation.transformation import (
     BaseImageGenerationConfig,
 )
@@ -407,7 +408,7 @@ class OpenRouterImageGenerationConfig(BaseImageGenerationConfig):
 
     def get_error_class(
         self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers]
-    ) -> Exception:
+    ) -> BaseLLMException:
         """Get the appropriate error class for OpenRouter errors."""
         return OpenRouterException(
             message=error_message,
diff --git a/tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py b/tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py
new file mode 100644
index 0000000000..a247b3c027
--- /dev/null
+++ b/tests/test_litellm/llms/openrouter/image_generation/test_openrouter_image_gen_transformation.py
@@ -0,0 +1,573 @@
+import json
+import os
+import sys
+from unittest.mock import MagicMock, patch
+
+import httpx
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../../../../..")
+)  # Adds the parent directory to the system path
+
+from litellm.llms.openrouter.image_generation.transformation import (
+    OpenRouterImageGenerationConfig,
+)
+from litellm.llms.openrouter.common_utils import OpenRouterException
+from litellm.types.utils import ImageResponse
+
+
+class TestOpenRouterImageGenerationTransformation:
+    def setup_method(self):
+        """Set up test fixtures before each test method."""
+        self.config = OpenRouterImageGenerationConfig()
+        self.model = "google/gemini-2.5-flash-image"
+        self.logging_obj = MagicMock()
+
+    def test_get_supported_openai_params(self):
+        """Test that get_supported_openai_params returns correct parameters."""
+        supported_params = self.config.get_supported_openai_params(self.model)
+        
+        assert "size" in supported_params
+        assert "quality" in supported_params
+        assert "n" in supported_params
+        assert len(supported_params) == 3
+
+    def test_map_size_to_aspect_ratio_square(self):
+        """Test mapping square sizes to aspect ratio."""
+        assert self.config._map_size_to_aspect_ratio("256x256") == "1:1"
+        assert self.config._map_size_to_aspect_ratio("512x512") == "1:1"
+        assert self.config._map_size_to_aspect_ratio("1024x1024") == "1:1"
+
+    def test_map_size_to_aspect_ratio_landscape(self):
+        """Test mapping landscape sizes to aspect ratio."""
+        assert self.config._map_size_to_aspect_ratio("1536x1024") == "3:2"
+        assert self.config._map_size_to_aspect_ratio("1792x1024") == "16:9"
+
+    def test_map_size_to_aspect_ratio_portrait(self):
+        """Test mapping portrait sizes to aspect ratio."""
+        assert self.config._map_size_to_aspect_ratio("1024x1536") == "2:3"
+        assert self.config._map_size_to_aspect_ratio("1024x1792") == "9:16"
+
+    def test_map_size_to_aspect_ratio_auto(self):
+        """Test mapping auto size to default aspect ratio."""
+        assert self.config._map_size_to_aspect_ratio("auto") == "1:1"
+
+    def test_map_size_to_aspect_ratio_unknown(self):
+        """Test mapping unknown size defaults to 1:1."""
+        assert self.config._map_size_to_aspect_ratio("999x999") == "1:1"
+
+    def test_map_quality_to_image_size_low(self):
+        """Test mapping low quality values to 1K."""
+        assert self.config._map_quality_to_image_size("low") == "1K"
+        assert self.config._map_quality_to_image_size("standard") == "1K"
+        assert self.config._map_quality_to_image_size("auto") == "1K"
+
+    def test_map_quality_to_image_size_medium(self):
+        """Test mapping medium quality to 2K."""
+        assert self.config._map_quality_to_image_size("medium") == "2K"
+
+    def test_map_quality_to_image_size_high(self):
+        """Test mapping high quality values to 4K."""
+        assert self.config._map_quality_to_image_size("high") == "4K"
+        assert self.config._map_quality_to_image_size("hd") == "4K"
+
+    def test_map_quality_to_image_size_unknown(self):
+        """Test mapping unknown quality returns None."""
+        assert self.config._map_quality_to_image_size("unknown") is None
+
+    def test_map_openai_params_size_only(self):
+        """Test that map_openai_params correctly maps size parameter."""
+        non_default_params = {"size": "1024x1024"}
+        optional_params = {}
+        
+        result = self.config.map_openai_params(
+            non_default_params=non_default_params,
+            optional_params=optional_params,
+            model=self.model,
+            drop_params=False
+        )
+        
+        assert "image_config" in result
+        assert result["image_config"]["aspect_ratio"] == "1:1"
+
+    def test_map_openai_params_quality_only(self):
+        """Test that map_openai_params correctly maps quality parameter."""
+        non_default_params = {"quality": "high"}
+        optional_params = {}
+        
+        result = self.config.map_openai_params(
+            non_default_params=non_default_params,
+            optional_params=optional_params,
+            model=self.model,
+            drop_params=False
+        )
+        
+        assert "image_config" in result
+        assert result["image_config"]["image_size"] == "4K"
+
+    def test_map_openai_params_size_and_quality(self):
+        """Test that map_openai_params correctly maps both size and quality."""
+        non_default_params = {
+            "size": "1792x1024",
+            "quality": "hd"
+        }
+        optional_params = {}
+        
+        result = self.config.map_openai_params(
+            non_default_params=non_default_params,
+            optional_params=optional_params,
+            model=self.model,
+            drop_params=False
+        )
+        
+        assert "image_config" in result
+        assert result["image_config"]["aspect_ratio"] == "16:9"
+        assert result["image_config"]["image_size"] == "4K"
+
+    def test_map_openai_params_with_n_parameter(self):
+        """Test that map_openai_params correctly passes through n parameter."""
+        non_default_params = {
+            "size": "1024x1024",
+            "n": 2
+        }
+        optional_params = {}
+        
+        result = self.config.map_openai_params(
+            non_default_params=non_default_params,
+            optional_params=optional_params,
+            model=self.model,
+            drop_params=False
+        )
+        
+        assert "image_config" in result
+        assert result["image_config"]["aspect_ratio"] == "1:1"
+        assert result["n"] == 2
+
+    def test_map_openai_params_unsupported_param_drop_false(self):
+        """Test that unsupported params are passed through when drop_params=False."""
+        non_default_params = {
+            "size": "1024x1024",
+            "unsupported_param": "value"
+        }
+        optional_params = {}
+        
+        result = self.config.map_openai_params(
+            non_default_params=non_default_params,
+            optional_params=optional_params,
+            model=self.model,
+            drop_params=False
+        )
+        
+        assert "image_config" in result
+        assert result["unsupported_param"] == "value"
+
+    def test_map_openai_params_unsupported_param_drop_true(self):
+        """Test that unsupported params are dropped when drop_params=True."""
+        non_default_params = {
+            "size": "1024x1024",
+            "unsupported_param": "value"
+        }
+        optional_params = {}
+        
+        result = self.config.map_openai_params(
+            non_default_params=non_default_params,
+            optional_params=optional_params,
+            model=self.model,
+            drop_params=True
+        )
+        
+        assert "image_config" in result
+        assert "unsupported_param" not in result
+
+    def test_get_complete_url_default(self):
+        """Test that get_complete_url returns default OpenRouter URL."""
+        result = self.config.get_complete_url(
+            api_base=None,
+            api_key="test_key",
+            model=self.model,
+            optional_params={},
+            litellm_params={}
+        )
+        
+        assert result == "https://openrouter.ai/api/v1/chat/completions"
+
+    def test_get_complete_url_with_custom_base(self):
+        """Test that get_complete_url uses custom api_base."""
+        custom_base = "https://custom.openrouter.ai/api/v1"
+        
+        result = self.config.get_complete_url(
+            api_base=custom_base,
+            api_key="test_key",
+            model=self.model,
+            optional_params={},
+            litellm_params={}
+        )
+        
+        assert result == f"{custom_base}/chat/completions"
+
+    def test_get_complete_url_with_base_already_complete(self):
+        """Test that get_complete_url doesn't duplicate /chat/completions."""
+        custom_base = "https://custom.openrouter.ai/api/v1/chat/completions"
+        
+        result = self.config.get_complete_url(
+            api_base=custom_base,
+            api_key="test_key",
+            model=self.model,
+            optional_params={},
+            litellm_params={}
+        )
+        
+        assert result == custom_base
+
+    @patch("litellm.llms.openrouter.image_generation.transformation.get_secret_str")
+    def test_validate_environment_with_api_key(self, mock_get_secret):
+        """Test that validate_environment correctly sets authorization header."""
+        headers = {}
+        api_key = "test_api_key"
+        
+        result = self.config.validate_environment(
+            headers=headers,
+            model=self.model,
+            messages=[],
+            optional_params={},
+            litellm_params={},
+            api_key=api_key
+        )
+        
+        assert result["Authorization"] == f"Bearer {api_key}"
+        mock_get_secret.assert_not_called()
+
+    @patch("litellm.llms.openrouter.image_generation.transformation.get_secret_str")
+    def test_validate_environment_with_secret_key(self, mock_get_secret):
+        """Test that validate_environment uses secret API key when api_key is None."""
+        mock_get_secret.return_value = "secret_api_key"
+        headers = {}
+        
+        result = self.config.validate_environment(
+            headers=headers,
+            model=self.model,
+            messages=[],
+            optional_params={},
+            litellm_params={},
+            api_key=None
+        )
+        
+        assert result["Authorization"] == "Bearer secret_api_key"
+        mock_get_secret.assert_called_once_with("OPENROUTER_API_KEY")
+
+    def test_transform_image_generation_request_basic(self):
+        """Test that transform_image_generation_request creates correct request body."""
+        prompt = "A beautiful sunset over mountains"
+        optional_params = {}
+        
+        result = self.config.transform_image_generation_request(
+            model=self.model,
+            prompt=prompt,
+            optional_params=optional_params,
+            litellm_params={},
+            headers={}
+        )
+        
+        assert result["model"] == self.model
+        assert result["messages"] == [{"role": "user", "content": prompt}]
+        assert "modalities" not in result  # modalities should not be added by default
+
+    def test_transform_image_generation_request_with_image_config(self):
+        """Test that transform_image_generation_request includes image_config."""
+        prompt = "A beautiful sunset"
+        optional_params = {
+            "image_config": {
+                "aspect_ratio": "16:9",
+                "image_size": "4K"
+            },
+            "n": 2
+        }
+        
+        result = self.config.transform_image_generation_request(
+            model=self.model,
+            prompt=prompt,
+            optional_params=optional_params,
+            litellm_params={},
+            headers={}
+        )
+        
+        assert result["model"] == self.model
+        assert result["messages"] == [{"role": "user", "content": prompt}]
+        assert result["image_config"]["aspect_ratio"] == "16:9"
+        assert result["image_config"]["image_size"] == "4K"
+        assert result["n"] == 2
+
+    def test_transform_image_generation_response_with_base64_images(self):
+        """Test that transform_image_generation_response correctly extracts base64 images."""
+        response_data = {
+            "choices": [{
+                "message": {
+                    "content": "Here is your image!",
+                    "role": "assistant",
+                    "images": [{
+                        "image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANS"},
+                        "index": 0,
+                        "type": "image_url"
+                    }]
+                }
+            }],
+            "usage": {
+                "prompt_tokens": 10,
+                "completion_tokens": 1300,
+                "total_tokens": 1310,
+                "completion_tokens_details": {"image_tokens": 1290},
+                "cost": 0.0387243
+            },
+            "model": "google/gemini-2.5-flash-image"
+        }
+        
+        mock_response = MagicMock()
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+        
+        model_response = ImageResponse(data=[])
+        
+        result = self.config.transform_image_generation_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=self.logging_obj,
+            request_data={},
+            optional_params={},
+            litellm_params={},
+            encoding=None
+        )
+        
+        assert len(result.data) == 1
+        assert result.data[0].b64_json == "iVBORw0KGgoAAAANS"
+        assert result.data[0].url is None
+
+    def test_transform_image_generation_response_with_url_images(self):
+        """Test that transform_image_generation_response correctly extracts URL images."""
+        response_data = {
+            "choices": [{
+                "message": {
+                    "content": "Here is your image!",
+                    "role": "assistant",
+                    "images": [{
+                        "image_url": {"url": "https://example.com/image.png"},
+                        "index": 0,
+                        "type": "image_url"
+                    }]
+                }
+            }],
+            "usage": {
+                "prompt_tokens": 10,
+                "completion_tokens": 1300,
+                "total_tokens": 1310
+            },
+            "model": "google/gemini-2.5-flash-image"
+        }
+        
+        mock_response = MagicMock()
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+        
+        model_response = ImageResponse(data=[])
+        
+        result = self.config.transform_image_generation_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=self.logging_obj,
+            request_data={},
+            optional_params={},
+            litellm_params={},
+            encoding=None
+        )
+        
+        assert len(result.data) == 1
+        assert result.data[0].url == "https://example.com/image.png"
+        assert result.data[0].b64_json is None
+
+    def test_transform_image_generation_response_with_usage_and_cost(self):
+        """Test that transform_image_generation_response correctly extracts usage and cost."""
+        response_data = {
+            "choices": [{
+                "message": {
+                    "content": "Here is your image!",
+                    "role": "assistant",
+                    "images": [{
+                        "image_url": {"url": "data:image/png;base64,abc123"},
+                        "index": 0,
+                        "type": "image_url"
+                    }]
+                }
+            }],
+            "usage": {
+                "prompt_tokens": 10,
+                "completion_tokens": 1300,
+                "total_tokens": 1310,
+                "completion_tokens_details": {"image_tokens": 1290},
+                "cost": 0.0387243,
+                "cost_details": {"input_cost": 0.001, "output_cost": 0.037}
+            },
+            "model": "google/gemini-2.5-flash-image"
+        }
+        
+        mock_response = MagicMock()
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+        
+        model_response = ImageResponse(data=[])
+        
+        result = self.config.transform_image_generation_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=self.logging_obj,
+            request_data={},
+            optional_params={},
+            litellm_params={},
+            encoding=None
+        )
+        
+        # Check usage
+        assert result.usage is not None
+        assert result.usage.input_tokens == 10
+        assert result.usage.output_tokens == 1290
+        assert result.usage.total_tokens == 1310
+        assert result.usage.input_tokens_details.text_tokens == 10
+        assert result.usage.input_tokens_details.image_tokens == 0
+        
+        # Check cost
+        assert hasattr(result, "_hidden_params")
+        assert "additional_headers" in result._hidden_params
+        assert result._hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] == 0.0387243
+        
+        # Check cost details
+        assert "response_cost_details" in result._hidden_params
+        assert result._hidden_params["response_cost_details"]["input_cost"] == 0.001
+        assert result._hidden_params["response_cost_details"]["output_cost"] == 0.037
+        
+        # Check model
+        assert result._hidden_params["model"] == "google/gemini-2.5-flash-image"
+
+    def test_transform_image_generation_response_multiple_images(self):
+        """Test that transform_image_generation_response handles multiple images."""
+        response_data = {
+            "choices": [{
+                "message": {
+                    "content": "Here are your images!",
+                    "role": "assistant",
+                    "images": [
+                        {
+                            "image_url": {"url": "data:image/png;base64,image1data"},
+                            "index": 0,
+                            "type": "image_url"
+                        },
+                        {
+                            "image_url": {"url": "data:image/png;base64,image2data"},
+                            "index": 1,
+                            "type": "image_url"
+                        }
+                    ]
+                }
+            }],
+            "usage": {
+                "prompt_tokens": 10,
+                "completion_tokens": 2600,
+                "total_tokens": 2610
+            },
+            "model": "google/gemini-2.5-flash-image"
+        }
+        
+        mock_response = MagicMock()
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+        
+        model_response = ImageResponse(data=[])
+        
+        result = self.config.transform_image_generation_response(
+            model=self.model,
+            raw_response=mock_response,
+            model_response=model_response,
+            logging_obj=self.logging_obj,
+            request_data={},
+            optional_params={},
+            litellm_params={},
+            encoding=None
+        )
+        
+        assert len(result.data) == 2
+        assert result.data[0].b64_json == "image1data"
+        assert result.data[1].b64_json == "image2data"
+
+    def test_transform_image_generation_response_json_error(self):
+        """Test that transform_image_generation_response raises error on invalid JSON."""
+        mock_response = MagicMock()
+        mock_response.json.side_effect = json.JSONDecodeError("Invalid JSON", "", 0)
+        mock_response.status_code = 500
+        mock_response.headers = {}
+        
+        model_response = ImageResponse(data=[])
+        
+        with pytest.raises(OpenRouterException) as exc_info:
+            self.config.transform_image_generation_response(
+                model=self.model,
+                raw_response=mock_response,
+                model_response=model_response,
+                logging_obj=self.logging_obj,
+                request_data={},
+                optional_params={},
+                litellm_params={},
+                encoding=None
+            )
+        
+        assert "Error parsing OpenRouter response" in str(exc_info.value)
+        assert exc_info.value.status_code == 500
+
+    def test_transform_image_generation_response_transformation_error(self):
+        """Test that transform_image_generation_response handles transformation errors."""
+        response_data = {
+            "choices": [{
+                "message": {
+                    "content": "Here is your image!",
+                    "role": "assistant",
+                    "images": "invalid_format"  # Invalid format
+                }
+            }]
+        }
+        
+        mock_response = MagicMock()
+        mock_response.json.return_value = response_data
+        mock_response.status_code = 200
+        mock_response.headers = {}
+        
+        model_response = ImageResponse(data=[])
+        
+        with pytest.raises(OpenRouterException) as exc_info:
+            self.config.transform_image_generation_response(
+                model=self.model,
+                raw_response=mock_response,
+                model_response=model_response,
+                logging_obj=self.logging_obj,
+                request_data={},
+                optional_params={},
+                litellm_params={},
+                encoding=None
+            )
+        
+        assert "Error transforming OpenRouter image generation response" in str(exc_info.value)
+
+    def test_get_error_class(self):
+        """Test that get_error_class returns OpenRouterException."""
+        error = self.config.get_error_class(
+            error_message="Test error",
+            status_code=400,
+            headers={"Content-Type": "application/json"}
+        )
+        
+        assert isinstance(error, OpenRouterException)
+        assert "Test error" in str(error)
+        assert error.status_code == 400