test(integration): databricks cached prompt tokens bill at cache rates (Pylon #7738)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-09-25 01:02:15 +00:00 · 2026-09-22 21:22:41 +00:00 · 2026-09-22 21:22:41 +00:00 · 22d86f4fd3
commit 22d86f4fd3
parent d567395609
1 changed files with 89 additions and 0 deletions
--- a/tests/integration/pricing/test_databricks_cache_pricing.py
+++ b/tests/integration/pricing/test_databricks_cache_pricing.py
@ -0,0 +1,89 @@
+import json
+import uuid
+from typing import Final
+
+import pytest
+
+from tests.integration._support.client import Gateway, eventually, object_value, string_value
+from tests.integration._support.database import read_rows
+from tests.integration._support.upstream import delete_scenario, register_scenario
+from tests.integration.cost_calculation.cost_tracking_case import JsonResponse
+
+INPUT_RATE: Final = 0.001
+OUTPUT_RATE: Final = 0.002
+CACHE_CREATION_RATE: Final = 0.004
+CACHE_READ_RATE: Final = 0.0001
+UNCACHED_PROMPT_TOKENS: Final = 1000
+CACHE_CREATION_TOKENS: Final = 2000
+CACHE_READ_TOKENS: Final = 8000
+PROMPT_TOKENS: Final = UNCACHED_PROMPT_TOKENS + CACHE_CREATION_TOKENS + CACHE_READ_TOKENS
+COMPLETION_TOKENS: Final = 500
+
+
+def databricks_cached_response() -> JsonResponse:
+    return JsonResponse(
+        content_type="application/json",
+        body={
+            "id": "chatcmpl-$REQUEST_ID",
+            "object": "chat.completion",
+            "created": 1700000000,
+            "model": "databricks-claude-integration",
+            "choices": [
+                {"index": 0, "message": {"role": "assistant", "content": "cached reply"}, "finish_reason": "stop"}
+            ],
+            "usage": {
+                "prompt_tokens": PROMPT_TOKENS,
+                "completion_tokens": COMPLETION_TOKENS,
+                "total_tokens": PROMPT_TOKENS + COMPLETION_TOKENS,
+                "cache_creation_input_tokens": CACHE_CREATION_TOKENS,
+                "cache_read_input_tokens": CACHE_READ_TOKENS,
+            },
+        },
+    )
+
+
+@pytest.mark.covers("pricing.databricks.cached_prompt_tokens_bill_at_cache_rates")
+def test_databricks_cached_prompt_tokens_bill_at_cache_rates_not_input_rate(gateway: Gateway) -> None:
+    with gateway.scenario() as scenario:
+        scenario_id: Final = f"databricks-cache-{uuid.uuid4().hex[:12]}"
+        handle: Final = register_scenario(scenario_id, databricks_cached_response())
+        scenario.cleanups.callback(delete_scenario, handle)
+        model: Final = scenario.model(
+            model="databricks/databricks-claude-integration",
+            api_base=handle.api_base(),
+            input_cost_per_token=INPUT_RATE,
+            output_cost_per_token=OUTPUT_RATE,
+            cache_creation_input_token_cost=CACHE_CREATION_RATE,
+            cache_read_input_token_cost=CACHE_READ_RATE,
+        )
+        response: Final = gateway.request(
+            "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": "cache control"}]}
+        )
+        assert response.status_code == 200, response.text
+        expected_prompt_cost: Final = (
+            UNCACHED_PROMPT_TOKENS * INPUT_RATE
+            + CACHE_CREATION_TOKENS * CACHE_CREATION_RATE
+            + CACHE_READ_TOKENS * CACHE_READ_RATE
+        )
+        expected_completion_cost: Final = COMPLETION_TOKENS * OUTPUT_RATE
+        assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(
+            expected_prompt_cost + expected_completion_cost, rel=1e-6
+        ), response.text
+        request_id: Final = string_value(object_value(response.json())["id"])
+        rows: Final = eventually(
+            lambda: read_rows(
+                'SELECT spend, metadata, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" '
+                "WHERE request_id = %s",
+                (request_id,),
+            ),
+            lambda values: len(values) == 1,
+            seconds=70,
+        )
+        assert rows[0]["prompt_tokens"] == PROMPT_TOKENS
+        assert rows[0]["completion_tokens"] == COMPLETION_TOKENS
+        assert float(rows[0]["spend"]) == pytest.approx(expected_prompt_cost + expected_completion_cost, rel=1e-6)
+        metadata: Final = rows[0]["metadata"]
+        parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata)
+        breakdown: Final = object_value(parsed["cost_breakdown"])
+        assert float(breakdown["input_cost"]) == pytest.approx(expected_prompt_cost, rel=1e-6)
+        assert float(breakdown["output_cost"]) == pytest.approx(expected_completion_cost, rel=1e-6)