From 9bd3463571417b241feefe352880956ff005e033 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 21:56:25 -0700 Subject: [PATCH 1/4] test(e2e): cover persisted pricing spend --- .../coverage_registry/quota_management.yaml | 3 + .../spend_tracking/cost_rows.py | 7 + .../test_pricing_config_spend_e2e.py | 247 ++++++++++++++++++ 3 files changed, 257 insertions(+) create mode 100644 tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index ad0914d455b..7cf4f562c4d 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/tests/e2e/coverage_registry/quota_management.yaml @@ -43,6 +43,9 @@ - {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/experimental_pass_through/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"} - {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"} - {id: quota_management.spend_tracking.cache_hit.zero_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: cache_hit, assertions: [zero_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A response-cache hit logs at zero cost with the cache-hit marker"} +- {id: quota_management.spend_tracking.discount_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: discount_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider discount changes the real request's spend row"} +- {id: quota_management.spend_tracking.margin_config.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: margin_config, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/management_endpoints/cost_tracking_settings.py + cost_calculator.py", rationale: "A persisted provider margin is added to the real request's spend row"} +- {id: quota_management.spend_tracking.guardrail_cost.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: guardrail_cost, assertions: [logs_cost], exercised_on: [chat_completions], source: "litellm_core_utils/llm_cost_calc/guardrail_cost.py + proxy/spend_tracking/spend_tracking_utils.py", rationale: "A priced Bedrock guardrail charge reaches the persisted spend row"} - {id: quota_management.spend_tracking.key_rollup.matches_sum_of_logs, module: quota_management, tier: P1, behavior: spend_tracking, variant: key_rollup, assertions: [matches_sum_of_logs], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "A key's rolled-up spend equals the sum of its log rows"} - {id: quota_management.spend_tracking.concurrent_burst.loses_no_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: concurrent_burst, assertions: [loses_no_spend], exercised_on: [chat_completions], source: "proxy/db/db_spend_update_writer.py", rationale: "Concurrent calls all land as spend; no row lost to write contention"} - {id: quota_management.spend_tracking.tags.attributes_spend, module: quota_management, tier: P1, behavior: spend_tracking, variant: tags, assertions: [attributes_spend], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Request tags round-trip to spend rows and tag rollups match tagged logs"} diff --git a/tests/e2e/quota_management/spend_tracking/cost_rows.py b/tests/e2e/quota_management/spend_tracking/cost_rows.py index 87af54fe83f..c58247ea555 100644 --- a/tests/e2e/quota_management/spend_tracking/cost_rows.py +++ b/tests/e2e/quota_management/spend_tracking/cost_rows.py @@ -49,6 +49,13 @@ class CostBreakdownRow(BaseModel): reasoning_cost: float | None = None tool_usage_cost: float | None = None total_cost: float | None = None + original_cost: float | None = None + discount_percent: float | None = None + discount_amount: float | None = None + margin_percent: float | None = None + margin_fixed_amount: float | None = None + margin_total_amount: float | None = None + guardrail_cost: float | None = None service_tier: str | None = None diff --git a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py new file mode 100644 index 00000000000..f3f1882bd53 --- /dev/null +++ b/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import time +from collections.abc import Iterator +from typing import Final, Literal + +import pytest +from cost_rows import CostRow, approx_equal, poll_cost_row, register_priced_model +from e2e_config import settle_propagation, unique_marker +from e2e_http import NoBody, StreamingResponse, unwrap +from lifecycle import ResourceManager +from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody +from pydantic import BaseModel, RootModel +from spend_e2e_client import SpendClient + +pytestmark = pytest.mark.e2e + +INPUT_RATE: Final = 0.00005 +OUTPUT_RATE: Final = 0.0001 +DISCOUNT: Final = 0.25 +MARGIN_PERCENT: Final = 0.1 +MARGIN_FIXED: Final = 0.0005 + + +class _ConfigPatchResponse(BaseModel): + status: str + values: dict[str, float | dict[str, float]] + + +class _DiscountConfig(RootModel[dict[str, float]]): + pass + + +class _MarginConfig(RootModel[dict[str, float | dict[str, float]]]): + pass + + +class _BedrockChecksCategory(BaseModel): + category: Literal["VIOLENCE", "HATE", "SEXUAL", "MISCONDUCT", "INSULTS"] + + +class _BedrockContentFilter(BaseModel): + categories: list[_BedrockChecksCategory] + + +class _BedrockChecks(BaseModel): + contentFilter: _BedrockContentFilter + + +class _BedrockParams(BaseModel): + guardrail: Literal["bedrock"] = "bedrock" + mode: Literal["pre_call"] = "pre_call" + default_on: bool = False + checks: _BedrockChecks + + +class _GuardrailSpec(BaseModel): + guardrail_name: str + litellm_params: _BedrockParams + + +class _GuardrailCreate(BaseModel): + guardrail: _GuardrailSpec + + +class _GuardrailCreateResponse(BaseModel): + guardrail_id: str + + +def _register_bedrock_guardrail(client: SpendClient, resources: ResourceManager, name: str) -> None: + guardrail_id: Final = unwrap( + client.proxy.transport.post( + "/guardrails", + headers=client.proxy.transport.master, + json=_GuardrailCreate( + guardrail=_GuardrailSpec( + guardrail_name=name, + litellm_params=_BedrockParams( + checks=_BedrockChecks( + contentFilter=_BedrockContentFilter( + categories=[_BedrockChecksCategory(category="HATE")] + ) + ) + ), + ) + ), + response_type=_GuardrailCreateResponse, + ) + ).guardrail_id + settle_propagation(time.monotonic()) + resources.defer(lambda: _delete_guardrail(client, guardrail_id)) + + +def _delete_guardrail(client: SpendClient, guardrail_id: str) -> None: + _ = client.proxy.transport.delete( + f"/guardrails/{guardrail_id}", + headers=client.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) + + +def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> StreamingResponse: + return client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"reply with a short greeting {unique_marker()}")], + max_tokens=16, + guardrails=[name], + ), + ) + + +def _set_discount(client: SpendClient, values: dict[str, float]) -> None: + unwrap( + client.proxy.transport.patch( + "/config/cost_discount_config", + headers=client.proxy.transport.master, + json=_DiscountConfig(values), + response_type=_ConfigPatchResponse, + ) + ) + + +def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> None: + unwrap( + client.proxy.transport.patch( + "/config/cost_margin_config", + headers=client.proxy.transport.master, + json=_MarginConfig(values), + response_type=_ConfigPatchResponse, + ) + ) + + +def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str: + return register_priced_model( + client.proxy, + resources, + prefix, + LiteLLMParamsBody( + model="openai/gpt-4o-mini", + api_key="os.environ/OPENAI_API_KEY", + input_cost_per_token=INPUT_RATE, + output_cost_per_token=OUTPUT_RATE, + ), + ) + + +def _base_cost(row: CostRow, prompt_tokens: int, completion_tokens: int) -> float: + assert prompt_tokens > 0 and completion_tokens > 0 + assert row.prompt_tokens == prompt_tokens and row.completion_tokens == completion_tokens + base_cost: Final = prompt_tokens * INPUT_RATE + completion_tokens * OUTPUT_RATE + assert row.breakdown.original_cost is not None + assert approx_equal(row.breakdown.original_cost, base_cost) + return base_cost + + +@pytest.fixture +def restored_pricing_config(client: SpendClient) -> Iterator[None]: + _set_discount(client, {}) + _set_margin(client, {}) + yield + _set_discount(client, {}) + _set_margin(client, {}) + + +class TestPricingConfigSpend: + @pytest.mark.covers( + "quota_management.spend_tracking.discount_config.logs_cost", + exercised_on=["chat_completions"], + ) + def test_configured_discount_reaches_persisted_spend_row( + self, + client: SpendClient, + resources: ResourceManager, + scoped_key: str, + restored_pricing_config: None, + ) -> None: + _set_discount(client, {"openai": DISCOUNT}) + model: Final = _register_model(client, resources, "discount-priced") + chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) + assert chat.id and chat.usage and chat.usage.prompt_tokens and chat.usage.completion_tokens + + row: Final = poll_cost_row(client.proxy, chat.id) + assert row is not None + base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) + breakdown: Final = row.breakdown + assert breakdown.discount_percent is not None and approx_equal(breakdown.discount_percent, DISCOUNT) + assert breakdown.discount_amount is not None and approx_equal(breakdown.discount_amount, base_cost * DISCOUNT) + assert row.spend is not None and approx_equal(row.spend, base_cost * (1 - DISCOUNT)) + + @pytest.mark.covers( + "quota_management.spend_tracking.margin_config.logs_cost", + exercised_on=["chat_completions"], + ) + def test_configured_margin_reaches_persisted_spend_row( + self, + client: SpendClient, + resources: ResourceManager, + scoped_key: str, + restored_pricing_config: None, + ) -> None: + _set_margin(client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}}) + model: Final = _register_model(client, resources, "margin-priced") + chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) + assert chat.id and chat.usage and chat.usage.prompt_tokens and chat.usage.completion_tokens + + row: Final = poll_cost_row(client.proxy, chat.id) + assert row is not None + base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) + breakdown: Final = row.breakdown + expected_margin: Final = base_cost * MARGIN_PERCENT + MARGIN_FIXED + assert breakdown.margin_percent is not None and approx_equal(breakdown.margin_percent, MARGIN_PERCENT) + assert breakdown.margin_fixed_amount is not None and approx_equal(breakdown.margin_fixed_amount, MARGIN_FIXED) + assert breakdown.margin_total_amount is not None and approx_equal(breakdown.margin_total_amount, expected_margin) + assert row.spend is not None and approx_equal(row.spend, base_cost + expected_margin) + + @pytest.mark.covers( + "quota_management.spend_tracking.guardrail_cost.logs_cost", + exercised_on=["chat_completions"], + ) + def test_bedrock_guardrail_cost_reaches_persisted_spend_row( + self, + client: SpendClient, + resources: ResourceManager, + scoped_key: str, + restored_pricing_config: None, + ) -> None: + name: Final = f"e2e-bedrock-cost-{unique_marker()}" + _register_bedrock_guardrail(client, resources, name) + model: Final = _register_model(client, resources, "guardrail-priced") + result: Final = _guarded_chat(client, scoped_key, model, name) + assert result.ok, f"guarded request failed with {result.status_code}: {result.body[:400]}" + assert name in {value.strip() for value in result.headers.get("x-litellm-applied-guardrails", "").split(",")} + + chat: Final = ChatResponse.model_validate_json(result.body) + assert chat.id and chat.usage and chat.usage.prompt_tokens and chat.usage.completion_tokens + row: Final = poll_cost_row(client.proxy, chat.id) + assert row is not None + base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) + breakdown: Final = row.breakdown + assert breakdown.guardrail_cost is not None and breakdown.guardrail_cost > 0 + assert breakdown.total_cost is not None and approx_equal(breakdown.total_cost, base_cost + breakdown.guardrail_cost) + assert row.spend is not None and approx_equal(row.spend, breakdown.total_cost) From 8538d9fe172feaed5cdb4912af04a9d0091aa74b Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 16 Sep 2026 22:59:02 -0700 Subject: [PATCH 2/4] test(e2e): isolate pricing spend state --- .../test_pricing_config_spend_e2e.py | 155 +++++++++++++----- 1 file changed, 117 insertions(+), 38 deletions(-) diff --git a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py index f3f1882bd53..a0df053b030 100644 --- a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py @@ -1,15 +1,24 @@ from __future__ import annotations +import os import time from collections.abc import Iterator from typing import Final, Literal import pytest -from cost_rows import CostRow, approx_equal, poll_cost_row, register_priced_model +from cost_rows import CostRow, approx_equal, poll_cost_row from e2e_config import settle_propagation, unique_marker from e2e_http import NoBody, StreamingResponse, unwrap from lifecycle import ResourceManager -from models import ChatBody, ChatMessage, ChatResponse, LiteLLMParamsBody +from models import ( + ChatBody, + ChatMessage, + ChatResponse, + KeyDeleteBody, + KeyGenerateBody, + LiteLLMParamsBody, + ModelDeleteBody, +) from pydantic import BaseModel, RootModel from spend_e2e_client import SpendClient @@ -35,23 +44,20 @@ class _MarginConfig(RootModel[dict[str, float | dict[str, float]]]): pass -class _BedrockChecksCategory(BaseModel): - category: Literal["VIOLENCE", "HATE", "SEXUAL", "MISCONDUCT", "INSULTS"] +class _DiscountConfigResponse(BaseModel): + values: dict[str, float] -class _BedrockContentFilter(BaseModel): - categories: list[_BedrockChecksCategory] - - -class _BedrockChecks(BaseModel): - contentFilter: _BedrockContentFilter +class _MarginConfigResponse(BaseModel): + values: dict[str, float | dict[str, float]] class _BedrockParams(BaseModel): guardrail: Literal["bedrock"] = "bedrock" mode: Literal["pre_call"] = "pre_call" default_on: bool = False - checks: _BedrockChecks + guardrailIdentifier: str + guardrailVersion: str class _GuardrailSpec(BaseModel): @@ -67,7 +73,13 @@ class _GuardrailCreateResponse(BaseModel): guardrail_id: str -def _register_bedrock_guardrail(client: SpendClient, resources: ResourceManager, name: str) -> None: +def _register_bedrock_guardrail( + client: SpendClient, + resources: ResourceManager, + name: str, + identifier: str, + version: str, +) -> None: guardrail_id: Final = unwrap( client.proxy.transport.post( "/guardrails", @@ -76,11 +88,8 @@ def _register_bedrock_guardrail(client: SpendClient, resources: ResourceManager, guardrail=_GuardrailSpec( guardrail_name=name, litellm_params=_BedrockParams( - checks=_BedrockChecks( - contentFilter=_BedrockContentFilter( - categories=[_BedrockChecksCategory(category="HATE")] - ) - ) + guardrailIdentifier=identifier, + guardrailVersion=version, ), ) ), @@ -92,11 +101,13 @@ def _register_bedrock_guardrail(client: SpendClient, resources: ResourceManager, def _delete_guardrail(client: SpendClient, guardrail_id: str) -> None: - _ = client.proxy.transport.delete( - f"/guardrails/{guardrail_id}", - headers=client.proxy.transport.master, - json=NoBody(), - response_type=NoBody, + unwrap( + client.proxy.transport.delete( + f"/guardrails/{guardrail_id}", + headers=client.proxy.transport.master, + json=NoBody(), + response_type=NoBody, + ) ) @@ -135,11 +146,32 @@ def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]] ) +def _get_discount(client: SpendClient) -> dict[str, float]: + return unwrap( + client.proxy.transport.get( + "/config/cost_discount_config", + headers=client.proxy.transport.master, + params=NoBody(), + response_type=_DiscountConfigResponse, + ) + ).values + + +def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]: + return unwrap( + client.proxy.transport.get( + "/config/cost_margin_config", + headers=client.proxy.transport.master, + params=NoBody(), + response_type=_MarginConfigResponse, + ) + ).values + + def _register_model(client: SpendClient, resources: ResourceManager, prefix: str) -> str: - return register_priced_model( - client.proxy, - resources, - prefix, + model: Final = f"{prefix}-{unique_marker()}" + model_id: Final = client.proxy.create_model( + model, LiteLLMParamsBody( model="openai/gpt-4o-mini", api_key="os.environ/OPENAI_API_KEY", @@ -147,6 +179,30 @@ def _register_model(client: SpendClient, resources: ResourceManager, prefix: str output_cost_per_token=OUTPUT_RATE, ), ) + resources.defer(lambda: _delete_model(client, model_id)) + return model + + +def _delete_model(client: SpendClient, model_id: str) -> None: + unwrap( + client.proxy.transport.post( + "/model/delete", + headers=client.proxy.transport.master, + json=ModelDeleteBody(id=model_id), + response_type=NoBody, + ) + ) + + +def _delete_key(client: SpendClient, key: str) -> None: + unwrap( + client.proxy.transport.post( + "/key/delete", + headers=client.proxy.transport.master, + json=KeyDeleteBody(keys=[key]), + response_type=NoBody, + ) + ) def _base_cost(row: CostRow, prompt_tokens: int, completion_tokens: int) -> float: @@ -158,13 +214,34 @@ def _base_cost(row: CostRow, prompt_tokens: int, completion_tokens: int) -> floa return base_cost +@pytest.fixture +def strict_resources(client: SpendClient) -> Iterator[ResourceManager]: + manager: Final = ResourceManager(client=client.proxy, strict_cleanup=True) + manager.init() + yield manager + manager.teardown() + + +@pytest.fixture +def scoped_key(client: SpendClient, strict_resources: ResourceManager) -> str: + key: Final = client.proxy.generate_key(KeyGenerateBody(user_id="e2e-test-user")) + strict_resources.defer(lambda: _delete_key(client, key)) + return key + + @pytest.fixture def restored_pricing_config(client: SpendClient) -> Iterator[None]: - _set_discount(client, {}) - _set_margin(client, {}) - yield - _set_discount(client, {}) - _set_margin(client, {}) + discount: Final = _get_discount(client) + margin: Final = _get_margin(client) + try: + _set_discount(client, {}) + _set_margin(client, {}) + yield + finally: + try: + _set_discount(client, discount) + finally: + _set_margin(client, margin) class TestPricingConfigSpend: @@ -175,12 +252,12 @@ class TestPricingConfigSpend: def test_configured_discount_reaches_persisted_spend_row( self, client: SpendClient, - resources: ResourceManager, + strict_resources: ResourceManager, scoped_key: str, restored_pricing_config: None, ) -> None: _set_discount(client, {"openai": DISCOUNT}) - model: Final = _register_model(client, resources, "discount-priced") + model: Final = _register_model(client, strict_resources, "discount-priced") chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) assert chat.id and chat.usage and chat.usage.prompt_tokens and chat.usage.completion_tokens @@ -199,12 +276,12 @@ class TestPricingConfigSpend: def test_configured_margin_reaches_persisted_spend_row( self, client: SpendClient, - resources: ResourceManager, + strict_resources: ResourceManager, scoped_key: str, restored_pricing_config: None, ) -> None: _set_margin(client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}}) - model: Final = _register_model(client, resources, "margin-priced") + model: Final = _register_model(client, strict_resources, "margin-priced") chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) assert chat.id and chat.usage and chat.usage.prompt_tokens and chat.usage.completion_tokens @@ -225,13 +302,15 @@ class TestPricingConfigSpend: def test_bedrock_guardrail_cost_reaches_persisted_spend_row( self, client: SpendClient, - resources: ResourceManager, + strict_resources: ResourceManager, scoped_key: str, restored_pricing_config: None, ) -> None: + identifier: Final = os.environ["BEDROCK_GUARDRAIL_IDENTIFIER"] + version: Final = os.environ["BEDROCK_GUARDRAIL_VERSION"] name: Final = f"e2e-bedrock-cost-{unique_marker()}" - _register_bedrock_guardrail(client, resources, name) - model: Final = _register_model(client, resources, "guardrail-priced") + _register_bedrock_guardrail(client, strict_resources, name, identifier, version) + model: Final = _register_model(client, strict_resources, "guardrail-priced") result: Final = _guarded_chat(client, scoped_key, model, name) assert result.ok, f"guarded request failed with {result.status_code}: {result.body[:400]}" assert name in {value.strip() for value in result.headers.get("x-litellm-applied-guardrails", "").split(",")} From 90775f4f768c6a63e2ada2df401aeb45d5c0c83a Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 00:04:01 -0700 Subject: [PATCH 3/4] test(e2e): synchronize pricing assertions --- .../test_pricing_config_spend_e2e.py | 34 +++++++++++++------ 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py index a0df053b030..ad50d03bbb3 100644 --- a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py @@ -124,7 +124,7 @@ def _guarded_chat(client: SpendClient, key: str, model: str, name: str) -> Strea ) -def _set_discount(client: SpendClient, values: dict[str, float]) -> None: +def _set_discount(client: SpendClient, values: dict[str, float]) -> float: unwrap( client.proxy.transport.patch( "/config/cost_discount_config", @@ -133,9 +133,10 @@ def _set_discount(client: SpendClient, values: dict[str, float]) -> None: response_type=_ConfigPatchResponse, ) ) + return time.monotonic() -def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> None: +def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]]) -> float: unwrap( client.proxy.transport.patch( "/config/cost_margin_config", @@ -144,6 +145,7 @@ def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]] response_type=_ConfigPatchResponse, ) ) + return time.monotonic() def _get_discount(client: SpendClient) -> dict[str, float]: @@ -206,7 +208,7 @@ def _delete_key(client: SpendClient, key: str) -> None: def _base_cost(row: CostRow, prompt_tokens: int, completion_tokens: int) -> float: - assert prompt_tokens > 0 and completion_tokens > 0 + assert prompt_tokens > 0 and completion_tokens >= 0 assert row.prompt_tokens == prompt_tokens and row.completion_tokens == completion_tokens base_cost: Final = prompt_tokens * INPUT_RATE + completion_tokens * OUTPUT_RATE assert row.breakdown.original_cost is not None @@ -235,13 +237,15 @@ def restored_pricing_config(client: SpendClient) -> Iterator[None]: margin: Final = _get_margin(client) try: _set_discount(client, {}) - _set_margin(client, {}) + reset_at: Final = _set_margin(client, {}) + settle_propagation(reset_at) yield finally: try: _set_discount(client, discount) finally: - _set_margin(client, margin) + restored_at: Final = _set_margin(client, margin) + settle_propagation(restored_at) class TestPricingConfigSpend: @@ -256,10 +260,13 @@ class TestPricingConfigSpend: scoped_key: str, restored_pricing_config: None, ) -> None: - _set_discount(client, {"openai": DISCOUNT}) + discount_set_at: Final = _set_discount(client, {"openai": DISCOUNT}) + settle_propagation(discount_set_at) model: Final = _register_model(client, strict_resources, "discount-priced") chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) - assert chat.id and chat.usage and chat.usage.prompt_tokens and chat.usage.completion_tokens + assert chat.id and chat.usage + assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 + assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 row: Final = poll_cost_row(client.proxy, chat.id) assert row is not None @@ -280,10 +287,15 @@ class TestPricingConfigSpend: scoped_key: str, restored_pricing_config: None, ) -> None: - _set_margin(client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}}) + margin_set_at: Final = _set_margin( + client, {"openai": {"percentage": MARGIN_PERCENT, "fixed_amount": MARGIN_FIXED}} + ) + settle_propagation(margin_set_at) model: Final = _register_model(client, strict_resources, "margin-priced") chat: Final = unwrap(client.chat(scoped_key, model, f"reply with one word {unique_marker()}", max_tokens=16)) - assert chat.id and chat.usage and chat.usage.prompt_tokens and chat.usage.completion_tokens + assert chat.id and chat.usage + assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 + assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 row: Final = poll_cost_row(client.proxy, chat.id) assert row is not None @@ -316,7 +328,9 @@ class TestPricingConfigSpend: assert name in {value.strip() for value in result.headers.get("x-litellm-applied-guardrails", "").split(",")} chat: Final = ChatResponse.model_validate_json(result.body) - assert chat.id and chat.usage and chat.usage.prompt_tokens and chat.usage.completion_tokens + assert chat.id and chat.usage + assert chat.usage.prompt_tokens is not None and chat.usage.prompt_tokens > 0 + assert chat.usage.completion_tokens is not None and chat.usage.completion_tokens >= 0 row: Final = poll_cost_row(client.proxy, chat.id) assert row is not None base_cost: Final = _base_cost(row, chat.usage.prompt_tokens, chat.usage.completion_tokens) From 3e4f1ea43936f163c97231f326072dbbabd9fed2 Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Thu, 17 Sep 2026 09:45:21 -0700 Subject: [PATCH 4/4] fix(e2e): share pricing request models --- tests/e2e/models.py | 42 +++++++++++ .../test_pricing_config_spend_e2e.py | 74 +++++-------------- 2 files changed, 62 insertions(+), 54 deletions(-) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 9f49c5974d0..79f23c8c665 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -905,6 +905,48 @@ class CostMap(RootModel[dict[str, CostMapEntry]]): pass +class ConfigPatchResponse(BaseModel): + status: str + values: dict[str, float | dict[str, float]] + + +class CostDiscountConfig(RootModel[dict[str, float]]): + pass + + +class CostMarginConfig(RootModel[dict[str, float | dict[str, float]]]): + pass + + +class CostDiscountConfigResponse(BaseModel): + values: dict[str, float] + + +class CostMarginConfigResponse(BaseModel): + values: dict[str, float | dict[str, float]] + + +class BedrockGuardrailParams(BaseModel): + guardrail: Literal["bedrock"] = "bedrock" + mode: Literal["pre_call"] = "pre_call" + default_on: bool = False + guardrailIdentifier: str + guardrailVersion: str + + +class BedrockGuardrailSpec(BaseModel): + guardrail_name: str + litellm_params: BedrockGuardrailParams + + +class GuardrailCreateBody(BaseModel): + guardrail: BedrockGuardrailSpec + + +class GuardrailCreateResponse(BaseModel): + guardrail_id: str + + class FileEntry(BaseModel): id: str diff --git a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py b/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py index ad50d03bbb3..228eb652ac9 100644 --- a/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py +++ b/tests/e2e/quota_management/spend_tracking/test_pricing_config_spend_e2e.py @@ -3,7 +3,7 @@ from __future__ import annotations import os import time from collections.abc import Iterator -from typing import Final, Literal +from typing import Final import pytest from cost_rows import CostRow, approx_equal, poll_cost_row @@ -11,15 +11,23 @@ from e2e_config import settle_propagation, unique_marker from e2e_http import NoBody, StreamingResponse, unwrap from lifecycle import ResourceManager from models import ( + BedrockGuardrailParams, + BedrockGuardrailSpec, ChatBody, ChatMessage, ChatResponse, + ConfigPatchResponse, + CostDiscountConfig, + CostDiscountConfigResponse, + CostMarginConfig, + CostMarginConfigResponse, + GuardrailCreateBody, + GuardrailCreateResponse, KeyDeleteBody, KeyGenerateBody, LiteLLMParamsBody, ModelDeleteBody, ) -from pydantic import BaseModel, RootModel from spend_e2e_client import SpendClient pytestmark = pytest.mark.e2e @@ -31,48 +39,6 @@ MARGIN_PERCENT: Final = 0.1 MARGIN_FIXED: Final = 0.0005 -class _ConfigPatchResponse(BaseModel): - status: str - values: dict[str, float | dict[str, float]] - - -class _DiscountConfig(RootModel[dict[str, float]]): - pass - - -class _MarginConfig(RootModel[dict[str, float | dict[str, float]]]): - pass - - -class _DiscountConfigResponse(BaseModel): - values: dict[str, float] - - -class _MarginConfigResponse(BaseModel): - values: dict[str, float | dict[str, float]] - - -class _BedrockParams(BaseModel): - guardrail: Literal["bedrock"] = "bedrock" - mode: Literal["pre_call"] = "pre_call" - default_on: bool = False - guardrailIdentifier: str - guardrailVersion: str - - -class _GuardrailSpec(BaseModel): - guardrail_name: str - litellm_params: _BedrockParams - - -class _GuardrailCreate(BaseModel): - guardrail: _GuardrailSpec - - -class _GuardrailCreateResponse(BaseModel): - guardrail_id: str - - def _register_bedrock_guardrail( client: SpendClient, resources: ResourceManager, @@ -84,16 +50,16 @@ def _register_bedrock_guardrail( client.proxy.transport.post( "/guardrails", headers=client.proxy.transport.master, - json=_GuardrailCreate( - guardrail=_GuardrailSpec( + json=GuardrailCreateBody( + guardrail=BedrockGuardrailSpec( guardrail_name=name, - litellm_params=_BedrockParams( + litellm_params=BedrockGuardrailParams( guardrailIdentifier=identifier, guardrailVersion=version, ), ) ), - response_type=_GuardrailCreateResponse, + response_type=GuardrailCreateResponse, ) ).guardrail_id settle_propagation(time.monotonic()) @@ -129,8 +95,8 @@ def _set_discount(client: SpendClient, values: dict[str, float]) -> float: client.proxy.transport.patch( "/config/cost_discount_config", headers=client.proxy.transport.master, - json=_DiscountConfig(values), - response_type=_ConfigPatchResponse, + json=CostDiscountConfig(values), + response_type=ConfigPatchResponse, ) ) return time.monotonic() @@ -141,8 +107,8 @@ def _set_margin(client: SpendClient, values: dict[str, float | dict[str, float]] client.proxy.transport.patch( "/config/cost_margin_config", headers=client.proxy.transport.master, - json=_MarginConfig(values), - response_type=_ConfigPatchResponse, + json=CostMarginConfig(values), + response_type=ConfigPatchResponse, ) ) return time.monotonic() @@ -154,7 +120,7 @@ def _get_discount(client: SpendClient) -> dict[str, float]: "/config/cost_discount_config", headers=client.proxy.transport.master, params=NoBody(), - response_type=_DiscountConfigResponse, + response_type=CostDiscountConfigResponse, ) ).values @@ -165,7 +131,7 @@ def _get_margin(client: SpendClient) -> dict[str, float | dict[str, float]]: "/config/cost_margin_config", headers=client.proxy.transport.master, params=NoBody(), - response_type=_MarginConfigResponse, + response_type=CostMarginConfigResponse, ) ).values