From 673233418c0e43cebdaf017f50c4bf605c432565 Mon Sep 17 00:00:00 2001
From: Kartik Chilkoti <chilkotikartik@gmail.com>
Date: Wed, 16 Sep 2026 14:54:39 +0530
Subject: [PATCH 1/2] test(cost): add unit tests for token usage aggregation
 and precision bounds [skip ci]

---
 tests/test_token_usage_and_cost_bounds.py | 43 +++++++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100644 tests/test_token_usage_and_cost_bounds.py

diff --git a/tests/test_token_usage_and_cost_bounds.py b/tests/test_token_usage_and_cost_bounds.py
new file mode 100644
index 00000000000..2c11284dd66
--- /dev/null
+++ b/tests/test_token_usage_and_cost_bounds.py
@@ -0,0 +1,43 @@
+import pytest
+from typing import Dict, Any
+
+def test_token_usage_aggregation():
+    responses = [
+        {"prompt_tokens": 120, "completion_tokens": 45, "total_tokens": 165},
+        {"prompt_tokens": 80, "completion_tokens": 30, "total_tokens": 110},
+        {"prompt_tokens": 200, "completion_tokens": 90, "total_tokens": 290}
+    ]
+    
+    total_prompt = sum(r["prompt_tokens"] for r in responses)
+    total_completion = sum(r["completion_tokens"] for r in responses)
+    total_combined = sum(r["total_tokens"] for r in responses)
+    
+    assert total_prompt == 400
+    assert total_completion == 165
+    assert total_combined == 565
+    assert total_prompt + total_completion == total_combined
+
+def test_cost_calculation_boundary_zero_tokens():
+    input_cost_per_token = 0.0000015
+    output_cost_per_token = 0.0000020
+    
+    prompt_tokens = 0
+    completion_tokens = 0
+    
+    cost = (prompt_tokens * input_cost_per_token) + (completion_tokens * output_cost_per_token)
+    assert cost == 0.0
+
+def test_cost_calculation_precision():
+    input_cost_per_million = 2.50
+    output_cost_per_million = 10.00
+    
+    prompt_tokens = 100_000
+    completion_tokens = 50_000
+    
+    input_cost = (prompt_tokens / 1_000_000) * input_cost_per_million
+    output_cost = (completion_tokens / 1_000_000) * output_cost_per_million
+    total_cost = round(input_cost + output_cost, 4)
+    
+    assert input_cost == 0.25
+    assert output_cost == 0.50
+    assert total_cost == 0.75
\ No newline at end of file

From a0e16b13842e31d91e38105df6328084ff93b43a Mon Sep 17 00:00:00 2001
From: Kartik Chilkoti <chilkotikartik@gmail.com>
Date: Wed, 16 Sep 2026 20:18:02 +0530
Subject: [PATCH 2/2] test(cost): test litellm completion_cost, Usage
 dataclass, and model_cost lookup

---
 tests/test_token_usage_and_cost_bounds.py | 59 ++++++++---------------
 1 file changed, 21 insertions(+), 38 deletions(-)

diff --git a/tests/test_token_usage_and_cost_bounds.py b/tests/test_token_usage_and_cost_bounds.py
index 2c11284dd66..4544f951a1d 100644
--- a/tests/test_token_usage_and_cost_bounds.py
+++ b/tests/test_token_usage_and_cost_bounds.py
@@ -1,43 +1,26 @@
 import pytest
-from typing import Dict, Any
+import litellm
+from litellm import completion_cost, Usage
 
-def test_token_usage_aggregation():
-    responses = [
-        {"prompt_tokens": 120, "completion_tokens": 45, "total_tokens": 165},
-        {"prompt_tokens": 80, "completion_tokens": 30, "total_tokens": 110},
-        {"prompt_tokens": 200, "completion_tokens": 90, "total_tokens": 290}
-    ]
-    
-    total_prompt = sum(r["prompt_tokens"] for r in responses)
-    total_completion = sum(r["completion_tokens"] for r in responses)
-    total_combined = sum(r["total_tokens"] for r in responses)
-    
-    assert total_prompt == 400
-    assert total_completion == 165
-    assert total_combined == 565
-    assert total_prompt + total_completion == total_combined
+def test_usage_dataclass_initialization():
+    usage = Usage(prompt_tokens=150, completion_tokens=50, total_tokens=200)
+    assert usage.prompt_tokens == 150
+    assert usage.completion_tokens == 50
+    assert usage.total_tokens == 200
 
-def test_cost_calculation_boundary_zero_tokens():
-    input_cost_per_token = 0.0000015
-    output_cost_per_token = 0.0000020
-    
-    prompt_tokens = 0
-    completion_tokens = 0
-    
-    cost = (prompt_tokens * input_cost_per_token) + (completion_tokens * output_cost_per_token)
+def test_completion_cost_with_zero_tokens():
+    mock_resp = {
+        "model": "gpt-3.5-turbo",
+        "usage": {
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "total_tokens": 0
+        }
+    }
+    cost = completion_cost(completion_response=mock_resp)
     assert cost == 0.0
 
-def test_cost_calculation_precision():
-    input_cost_per_million = 2.50
-    output_cost_per_million = 10.00
-    
-    prompt_tokens = 100_000
-    completion_tokens = 50_000
-    
-    input_cost = (prompt_tokens / 1_000_000) * input_cost_per_million
-    output_cost = (completion_tokens / 1_000_000) * output_cost_per_million
-    total_cost = round(input_cost + output_cost, 4)
-    
-    assert input_cost == 0.25
-    assert output_cost == 0.50
-    assert total_cost == 0.75
\ No newline at end of file
+def test_model_cost_lookup():
+    cost_map = litellm.model_cost
+    assert isinstance(cost_map, dict)
+    assert "gpt-4" in cost_map or "gpt-3.5-turbo" in cost_map
\ No newline at end of file