diff --git a/tests/integration/_support/upstream.py b/tests/integration/_support/upstream.py index 3c4198c8133..a289589b2dc 100644 --- a/tests/integration/_support/upstream.py +++ b/tests/integration/_support/upstream.py @@ -1,19 +1,25 @@ from __future__ import annotations import argparse -import json -import os -import struct -import zlib from collections import deque from collections.abc import Mapping +import json from dataclasses import dataclass, field +import os from pathlib import Path from queue import SimpleQueue +import struct from typing import Final, cast +import zlib import httpx import uvicorn +from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.responses import JSONResponse, Response +from starlette.routing import Route + from _fake_openai_endpoint_server import chat_completions, completions, embeddings, health, moderations from integration.cost_calculation.cost_tracking_case import ( EventStreamResponse, @@ -21,11 +27,6 @@ from integration.cost_calculation.cost_tracking_case import ( SseResponse, StoredResponse, ) -from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError -from starlette.applications import Starlette -from starlette.requests import Request -from starlette.responses import JSONResponse, Response -from starlette.routing import Route JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) CASES_FILE: Final = Path(__file__).resolve().parents[1] / "cost_calculation" / "cost_tracking_cases.json" diff --git a/tests/integration/cost_calculation/conftest.py b/tests/integration/cost_calculation/conftest.py index 192f2bf5461..166488e36a5 100644 --- a/tests/integration/cost_calculation/conftest.py +++ b/tests/integration/cost_calculation/conftest.py @@ -9,11 +9,12 @@ from typing import Final from cryptography.hazmat.primitives import serialization from cryptography.hazmat.primitives.asymmetric import rsa +from pydantic import BaseModel, ConfigDict + from integration._support.client import JSON_OBJECT, Scenario, eventually, object_value, string_value from integration._support.database import read_rows from integration._support.upstream import delete_scenario, register_scenario from integration.cost_calculation.cost_tracking_case import CostTrackingTestCase -from pydantic import BaseModel, ConfigDict class CostBreakdown(BaseModel): diff --git a/tests/integration/cost_calculation/cost_tracking_cases.json b/tests/integration/cost_calculation/cost_tracking_cases.json index 379250d3123..428124ac4a9 100644 --- a/tests/integration/cost_calculation/cost_tracking_cases.json +++ b/tests/integration/cost_calculation/cost_tracking_cases.json @@ -21817,8 +21817,10 @@ "response": { "content_type": "application/json", "body": { + "id": "resp_$REQUEST_ID", "object": "response", "status": "completed", + "created_at": 1700000000, "model": "gpt-5.6", "output": [ { diff --git a/tests/integration/cost_calculation/test_cost_tracking.py b/tests/integration/cost_calculation/test_cost_tracking.py index c9dc7a5ffe3..5b876974c63 100644 --- a/tests/integration/cost_calculation/test_cost_tracking.py +++ b/tests/integration/cost_calculation/test_cost_tracking.py @@ -6,6 +6,7 @@ from hashlib import sha256 from typing import Final, cast import pytest + from integration._support.client import JSON_OBJECT, Gateway from integration.cost_calculation.conftest import ( approx_equal, @@ -105,15 +106,34 @@ def test_case_bills_expected_cost(gateway: Gateway, case: CostTrackingTestCase) assert breakdown.output_cost is not None and approx_equal(breakdown.output_cost, expected.output_cost), ( f"{case.name}: output_cost {breakdown.output_cost} != expected {expected.output_cost}" ) - for field, header_name, expected_component in ( - ("cache_read_cost", "x-litellm-response-cost-cache-read", expected.cache_read_cost), - ("cache_creation_cost", "x-litellm-response-cost-cache-creation", expected.cache_creation_cost), - ("reasoning_cost", "x-litellm-response-cost-reasoning", expected.reasoning_cost), - ("tool_usage_cost", "x-litellm-response-cost-tool-usage", expected.tool_usage_cost), + for field, header_name, actual_component, expected_component in ( + ( + "cache_read_cost", + "x-litellm-response-cost-cache-read", + breakdown.cache_read_cost, + expected.cache_read_cost, + ), + ( + "cache_creation_cost", + "x-litellm-response-cost-cache-creation", + breakdown.cache_creation_cost, + expected.cache_creation_cost, + ), + ( + "reasoning_cost", + "x-litellm-response-cost-reasoning", + breakdown.reasoning_cost, + expected.reasoning_cost, + ), + ( + "tool_usage_cost", + "x-litellm-response-cost-tool-usage", + breakdown.tool_usage_cost, + expected.tool_usage_cost, + ), ): if expected_component is None: continue - actual_component: Final = getattr(breakdown, field) assert actual_component is not None and approx_equal(actual_component, expected_component), ( f"{case.name}: {field} {actual_component} != expected {expected_component}" )