From 7d1ea76130c4bbdd7f8f7521b2ef336e427bd26f Mon Sep 17 00:00:00 2001
From: Robert Keus <robertkeus@gmail.com>
Date: Mon, 27 Jul 2026 19:12:10 +0200
Subject: [PATCH] test(greenpt): cover public rerank cost path

---
 .../test_greenpt_rerank_transformation.py     | 28 +++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/tests/test_litellm/llms/greenpt/rerank/test_greenpt_rerank_transformation.py b/tests/test_litellm/llms/greenpt/rerank/test_greenpt_rerank_transformation.py
index 0902acd09e8..5eb6c7c0c7a 100644
--- a/tests/test_litellm/llms/greenpt/rerank/test_greenpt_rerank_transformation.py
+++ b/tests/test_litellm/llms/greenpt/rerank/test_greenpt_rerank_transformation.py
@@ -1,8 +1,26 @@
+import json
+from pathlib import Path
+
 import pytest
 
+import litellm
+from litellm.cost_calculator import rerank_cost
 from litellm.llms.greenpt.rerank.transformation import GreenPTRerankConfig
 
 
+@pytest.fixture
+def local_model_cost_map():
+    original_model_cost = litellm.model_cost
+    model_cost_path = Path(__file__).resolve().parents[5] / "model_prices_and_context_window.json"
+    litellm.model_cost = json.loads(model_cost_path.read_text(encoding="utf-8"))
+    litellm.get_model_info.cache_clear()
+    try:
+        yield
+    finally:
+        litellm.model_cost = original_model_cost
+        litellm.get_model_info.cache_clear()
+
+
 def test_greenpt_supported_params():
     assert GreenPTRerankConfig().get_supported_cohere_rerank_params("green-rerank") == [
         "query",
@@ -76,6 +94,16 @@ def test_calculate_rerank_cost():
     assert completion_cost == 0.0
 
 
+def test_public_rerank_cost_uses_token_pricing(local_model_cost_map):
+    prompt_cost, completion_cost = rerank_cost(
+        model="greenpt/green-rerank",
+        custom_llm_provider=None,
+        billed_units={"total_tokens": 1000},
+    )
+    assert prompt_cost == pytest.approx(0.000136524)
+    assert completion_cost == 0.0
+
+
 @pytest.mark.parametrize(
     ("billed_units", "model_info"),
     [