mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-02 02:11:58 +00:00
test: scope local cost map env var with monkeypatch to avoid test pollution
This commit is contained in:
parent
ca834e7bfc
commit
aaf40f398d
1 changed files with 12 additions and 14 deletions
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue