99 lines
3.6 KiB
Python
99 lines
3.6 KiB
Python
from contextlib import contextmanager
|
||
from datetime import datetime, timezone
|
||
from decimal import Decimal
|
||
from pathlib import Path
|
||
from types import SimpleNamespace
|
||
import unittest
|
||
from unittest.mock import patch
|
||
from uuid import uuid4
|
||
|
||
from core.capabilities import ModelCapabilities
|
||
from core.pricing import calculate_chat_cost, resolve_chat_price
|
||
from core.storage.usage import record_chat_usage
|
||
|
||
|
||
MODELS_DIR = Path(__file__).resolve().parents[1] / "config" / "models"
|
||
|
||
|
||
class PricingTests(unittest.TestCase):
|
||
def test_deepseek_flash_selects_peak_and_off_peak(self) -> None:
|
||
caps = ModelCapabilities.load("deepseek_v4.flash", MODELS_DIR)
|
||
# 2026-08-17 是周一；01:00 UTC 为峰时起点，04:00 已回到谷时。
|
||
peak = resolve_chat_price(
|
||
caps.pricing,
|
||
occurred_at=datetime(2026, 8, 17, 1, 0, tzinfo=timezone.utc),
|
||
)
|
||
off_peak = resolve_chat_price(
|
||
caps.pricing,
|
||
occurred_at=datetime(2026, 8, 17, 4, 0, tzinfo=timezone.utc),
|
||
)
|
||
self.assertIsNotNone(peak)
|
||
self.assertIsNotNone(off_peak)
|
||
assert peak is not None and off_peak is not None
|
||
self.assertEqual(peak.tier, "peak")
|
||
self.assertEqual(peak.input_per_mtoken, Decimal("0.44"))
|
||
self.assertEqual(off_peak.tier, "default")
|
||
self.assertEqual(off_peak.output_per_mtoken, Decimal("0.66"))
|
||
|
||
def test_deepseek_weekend_is_off_peak(self) -> None:
|
||
caps = ModelCapabilities.load("deepseek_v4.pro", MODELS_DIR)
|
||
quote = resolve_chat_price(
|
||
caps.pricing,
|
||
occurred_at=datetime(2026, 8, 22, 2, 0, tzinfo=timezone.utc),
|
||
)
|
||
assert quote is not None
|
||
self.assertEqual(quote.tier, "default")
|
||
|
||
def test_usd_quote_calculates_cache_split_in_cny(self) -> None:
|
||
caps = ModelCapabilities.load("deepseek_v4.flash", MODELS_DIR)
|
||
quote = resolve_chat_price(
|
||
caps.pricing,
|
||
occurred_at=datetime(2026, 8, 17, 1, 0, tzinfo=timezone.utc),
|
||
)
|
||
assert quote is not None
|
||
result = calculate_chat_cost(
|
||
quote,
|
||
prompt_tokens=1_000_000,
|
||
completion_tokens=500_000,
|
||
cache_hit_tokens=800_000,
|
||
)
|
||
# miss:0.2*0.44*7.2 + hit:0.8*0.014*7.2 + out:0.5*1.32*7.2
|
||
self.assertEqual(result.total_cny, Decimal("5.466240"))
|
||
self.assertEqual(result.cache_miss_tokens, 200_000)
|
||
|
||
def test_local_catalog_precedes_litellm_cost_map(self) -> None:
|
||
caps = ModelCapabilities.load("deepseek_v4.flash", MODELS_DIR)
|
||
added = []
|
||
|
||
class FakeSession:
|
||
def add(self, row):
|
||
added.append(row)
|
||
|
||
@contextmanager
|
||
def fake_scope():
|
||
yield FakeSession()
|
||
|
||
with (
|
||
patch("core.storage.usage.session_scope", fake_scope),
|
||
patch("core.storage.usage._safe_chat_cost_usd", return_value=Decimal("999")),
|
||
):
|
||
cost = record_chat_usage(
|
||
task_id=None,
|
||
user_id=uuid4(),
|
||
message_id=None,
|
||
model_profile="deepseek_v4.flash",
|
||
prompt_tokens=1_000_000,
|
||
completion_tokens=0,
|
||
pricing=caps.pricing,
|
||
occurred_at=datetime(2026, 8, 17, 4, 0, tzinfo=timezone.utc),
|
||
response=SimpleNamespace(created=None),
|
||
)
|
||
|
||
self.assertEqual(cost, Decimal("1.584000"))
|
||
self.assertEqual(added[0].units["pricing_revision"], "deepseek-v4-20260816")
|
||
self.assertEqual(added[0].units["price_tier"], "default")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|