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()