test: scope local cost map env var with monkeypatch to avoid test pollution

This commit is contained in:
Srivatsa03 2026-06-04 13:51:25 -05:00
parent ca834e7bfc
commit aaf40f398d

View file

@ -7,28 +7,26 @@ AWS lists Titan Text Embeddings V2 at $0.02 per 1M input tokens
(= $0.00002 per 1K tokens = 2e-08 per token).
"""
import os
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" # Load from local file
import importlib
import litellm.litellm_core_utils.get_model_cost_map
import litellm
# Reload modules to pick up environment variable
importlib.reload(litellm.litellm_core_utils.get_model_cost_map)
importlib.reload(litellm)
class TestBedrockEmbeddingPricing:
"""Test suite for Bedrock embedding model pricing in the cost map."""
def test_titan_embed_v2_commercial_input_cost(self):
def test_titan_embed_v2_commercial_input_cost(self, monkeypatch):
"""Titan Text Embeddings V2 should be priced at $0.02 / 1M tokens (2e-08)."""
from litellm import model_cost
# Scope the local-cost-map flag to this test only, so it does not leak
# into sibling tests. monkeypatch restores the environment on teardown.
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
model = model_cost["amazon.titan-embed-text-v2:0"]
import litellm.litellm_core_utils.get_model_cost_map
import litellm
# Reload so the cost map is re-read from the local file with the flag set.
importlib.reload(litellm.litellm_core_utils.get_model_cost_map)
importlib.reload(litellm)
model = litellm.model_cost["amazon.titan-embed-text-v2:0"]
assert model["input_cost_per_token"] == 2e-08
assert model["output_cost_per_token"] == 0.0