zcbot/tests/test_usage_accounting.py

39 lines
1.2 KiB
Python

from decimal import Decimal
import unittest
from core.loop import _extract_usage_details
from core.storage.usage import _fallback_chat_cost_cny
class UsageAccountingTests(unittest.TestCase):
def test_extract_usage_details_includes_cache_tokens(self) -> None:
usage = {
"prompt_tokens": 1200,
"completion_tokens": 80,
"prompt_cache_hit_tokens": 900,
"prompt_cache_miss_tokens": 300,
"completion_tokens_details": {"reasoning_tokens": 12},
}
details = _extract_usage_details(usage)
self.assertEqual(details["tokens_in"], 1200)
self.assertEqual(details["tokens_out"], 80)
self.assertEqual(details["cache_hit_tokens"], 900)
self.assertEqual(details["cache_miss_tokens"], 300)
self.assertEqual(details["reasoning_tokens"], 12)
def test_fallback_chat_cost_uses_price_snapshots(self) -> None:
cost = _fallback_chat_cost_cny(
prompt_tokens=1_000_000,
completion_tokens=500_000,
input_cny_per_mtoken=1.0,
output_cny_per_mtoken=10.0,
)
self.assertEqual(cost, Decimal("6.000000"))
if __name__ == "__main__":
unittest.main()