From 932e615e4643488b315f4e967b44cfbdd8f56004 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kevin=28=EC=9D=B4=EC=A3=BC=EC=98=A4=29?= <xop4p@naver.com>
Date: Tue, 17 Mar 2026 11:49:26 +0900
Subject: [PATCH] test: add tests for custom_pricing vs provider response cost
 priority

---
 tests/test_litellm/test_cost_calculator.py | 149 +++++++++++++++++++++
 1 file changed, 149 insertions(+)

diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py
index 8f5c3ece0ca..e49a3327e92 100644
--- a/tests/test_litellm/test_cost_calculator.py
+++ b/tests/test_litellm/test_cost_calculator.py
@@ -67,6 +67,155 @@ def test_cost_calculator_with_response_cost_in_additional_headers():
     assert result == 1000
 
 
+def test_custom_pricing_skips_provider_response_cost():
+    """
+    When custom_pricing=True and provider response cost exists in hidden_params,
+    response_cost_calculator should ignore the provider cost and use custom pricing
+    via completion_cost() instead.
+    """
+    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+    litellm.model_cost = litellm.get_model_cost_map(url="")
+
+    # Register custom pricing in model_cost
+    custom_input_cost = 0.00015  # KRW per token
+    custom_output_cost = 0.0006  # KRW per token
+    litellm.register_model(
+        {
+            "openrouter/openai/gpt-4.1-nano": {
+                "input_cost_per_token": custom_input_cost,
+                "output_cost_per_token": custom_output_cost,
+                "litellm_provider": "openrouter",
+            }
+        }
+    )
+
+    prompt_tokens = 12
+    completion_tokens = 3
+    provider_usd_cost = 0.0000024  # USD cost from OpenRouter
+
+    response = ModelResponse(
+        id="test-id",
+        model="openrouter/openai/gpt-4.1-nano",
+        choices=[],
+        usage=Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+        ),
+    )
+    # Simulate OpenRouter setting provider response cost in hidden_params
+    response._hidden_params["additional_headers"] = {
+        "llm_provider-x-litellm-response-cost": provider_usd_cost
+    }
+
+    result = response_cost_calculator(
+        response_object=response,
+        model="openrouter/openai/gpt-4.1-nano",
+        custom_llm_provider="openrouter",
+        call_type="acompletion",
+        optional_params={},
+        cache_hit=None,
+        base_model=None,
+        custom_pricing=True,
+    )
+
+    expected_custom_cost = (
+        prompt_tokens * custom_input_cost + completion_tokens * custom_output_cost
+    )
+
+    # Must use custom pricing, NOT provider response cost
+    assert result != provider_usd_cost, (
+        f"Should not use provider response cost ({provider_usd_cost})"
+    )
+    assert result == pytest.approx(expected_custom_cost), (
+        f"Got {result}, expected {expected_custom_cost} from custom pricing"
+    )
+
+
+def test_provider_response_cost_used_when_no_custom_pricing():
+    """
+    When custom_pricing is False/None and provider response cost exists,
+    response_cost_calculator should return the provider cost (existing behavior).
+    """
+
+    class MockResponse(BaseModel):
+        _hidden_params = {
+            "additional_headers": {
+                "llm_provider-x-litellm-response-cost": 0.0000024
+            }
+        }
+
+    result = response_cost_calculator(
+        response_object=MockResponse(),
+        model="openrouter/openai/gpt-4.1-nano",
+        custom_llm_provider="openrouter",
+        call_type="acompletion",
+        optional_params={},
+        cache_hit=None,
+        base_model=None,
+        custom_pricing=False,
+    )
+
+    assert result == 0.0000024, (
+        f"Should use provider response cost, got {result}"
+    )
+
+
+def test_custom_pricing_without_provider_response_cost():
+    """
+    When custom_pricing=True and no provider response cost exists,
+    response_cost_calculator should use custom pricing via completion_cost().
+    This verifies existing behavior is preserved.
+    """
+    os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+    litellm.model_cost = litellm.get_model_cost_map(url="")
+
+    custom_input_cost = 0.001
+    custom_output_cost = 0.002
+    litellm.register_model(
+        {
+            "deepinfra/meta-llama/Llama-3.2-3B-Instruct": {
+                "input_cost_per_token": custom_input_cost,
+                "output_cost_per_token": custom_output_cost,
+                "litellm_provider": "deepinfra",
+            }
+        }
+    )
+
+    prompt_tokens = 10
+    completion_tokens = 5
+
+    response = ModelResponse(
+        id="test-id",
+        model="deepinfra/meta-llama/Llama-3.2-3B-Instruct",
+        choices=[],
+        usage=Usage(
+            prompt_tokens=prompt_tokens,
+            completion_tokens=completion_tokens,
+            total_tokens=prompt_tokens + completion_tokens,
+        ),
+    )
+    # No provider response cost in hidden_params
+
+    result = response_cost_calculator(
+        response_object=response,
+        model="deepinfra/meta-llama/Llama-3.2-3B-Instruct",
+        custom_llm_provider="deepinfra",
+        call_type="acompletion",
+        optional_params={},
+        cache_hit=None,
+        base_model=None,
+        custom_pricing=True,
+    )
+
+    expected_cost = (
+        prompt_tokens * custom_input_cost + completion_tokens * custom_output_cost
+    )
+    assert result == pytest.approx(expected_cost), (
+        f"Got {result}, expected {expected_cost} from custom pricing"
+    )
+
+
 def test_cost_calculator_with_usage(monkeypatch):
     os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
     litellm.model_cost = litellm.get_model_cost_map(url="")