diff --git a/tests/llm_translation/test_bedrock_embedding_pricing.py b/tests/llm_translation/test_bedrock_embedding_pricing.py index ce166928b19..099d73fed87 100644 --- a/tests/llm_translation/test_bedrock_embedding_pricing.py +++ b/tests/llm_translation/test_bedrock_embedding_pricing.py @@ -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