# -*- coding: utf-8 -*- """Unit tests: AI model catalog, price lookup and cost calculation (wcs/ai/model_catalog.py, UsageTracker.py).""" import datetime import os import re import sys import unittest HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, HERE) import stubs # noqa: E402 stubs.install() from wcs.ai import UsageTracker as tracker # noqa: E402 from wcs.ai import model_catalog as catalog # noqa: E402 SETTINGS_XML = os.path.join(stubs.ADDON_DIR, 'resources', 'settings.xml') class CostTests(unittest.TestCase): def cost(self, provider, model, usage): return tracker.calculate_cost(provider, model, usage) def test_openai_cached_tokens_are_not_billed_twice(self): result = self.cost('openai', 'gpt-6-luna', {'input_tokens': 1000, 'output_tokens': 500, 'input_tokens_details': {'cached_tokens': 400}}) self.assertAlmostEqual(result['input_cost'], 600 * 0.10 / 1e6) self.assertAlmostEqual(result['cached_input_cost'], 400 * 0.01 / 1e6) self.assertAlmostEqual(result['output_cost'], 500 * 0.50 / 1e6) self.assertAlmostEqual(result['cache_savings'], 400 * (0.10 - 0.01) / 1e6) self.assertEqual((result['input_tokens'], result['cached_tokens'], result['output_tokens']), (600, 400, 500)) def test_anthropic_cache_read_and_write(self): usage = {'input_tokens': 1000, 'cached_input_tokens': 2000, 'cache_creation_input_tokens': 500, 'output_tokens': 300} for raw in (usage, {'input_tokens': 1000, 'cache_read_input_tokens': 2000, 'cache_creation_input_tokens': 500, 'output_tokens': 300}): # our usage and the raw response result = self.cost('anthropic', 'claude-opus-5-5', raw) self.assertAlmostEqual(result['input_cost'], (1000 * 4.00 + 500 * 5.00) / 1e6) self.assertAlmostEqual(result['cached_input_cost'], 2000 * 0.20 / 1e6) self.assertAlmostEqual(result['total_cost'], (1000 * 4.00 + 500 * 5.00 + 2000 * 0.20 + 300 * 20.00) / 1e6) self.assertEqual(result['input_tokens'], 1500) def test_haiku_55_is_billed_by_prompt_length_and_sonnet_55_reads_cache_at_five_percent(self): short = self.cost('anthropic', 'claude-haiku-5-5', {'input_tokens': 1000, 'cached_input_tokens': 2000, 'cache_creation_input_tokens': 500, 'output_tokens': 300}) self.assertAlmostEqual(short['total_cost'], (1000 * 0.10 + 500 * 0.125 + 2000 * 0.01 + 300 * 0.50) / 1e6) long = self.cost('anthropic', 'claude-haiku-5-5', {'input_tokens': 90_000, 'cached_input_tokens': 20_000, 'output_tokens': 1000}) # 110K prompt self.assertAlmostEqual(long['total_cost'], (90_000 * 0.50 + 20_000 * 0.05 + 1000 * 2.50) / 1e6) sonnet = self.cost('anthropic', 'claude-sonnet-5-5', {'input_tokens': 0, 'cached_input_tokens': 1_000_000}) self.assertAlmostEqual(sonnet['total_cost'], 0.10) def test_the_provider_default_model_is_the_first_one_and_haiku_needs_no_fallbacks(self): self.assertEqual(catalog.DEFAULT_MODEL['anthropic'], 'claude-haiku-5-5') self.assertEqual(catalog.DEFAULT_MODEL['openai'], 'gpt-6-luna') self.assertIn('claude-haiku-5-5', catalog.CLAUDE_EFFORT_MODELS) self.assertNotIn('claude-haiku-5-5', catalog.CLAUDE_FALLBACK_MODELS) self.assertTrue(set(catalog.CLAUDE_FALLBACK_MODELS) <= set(catalog.CLAUDE_EFFORT_MODELS) <= set(catalog.MODELS['anthropic'])) def test_google_thinking_is_output_and_long_prompt_uses_the_long_tier(self): usage = tracker.usage_from_google({'promptTokenCount': 250_000, 'candidatesTokenCount': 600, 'thoughtsTokenCount': 400, 'cachedContentTokenCount': 100_000}) self.assertEqual(usage['completion_tokens'], 1000) result = self.cost('google', 'gemini-3.1-pro-preview', usage) self.assertAlmostEqual(result['total_cost'], (150_000 * 4.00 + 100_000 * 0.40 + 1000 * 18.00) / 1e6) short = self.cost('google', 'gemini-3.1-pro-preview', {'prompt_tokens': 1000, 'completion_tokens': 10}) self.assertAlmostEqual(short['total_cost'], (1000 * 2.00 + 10 * 12.00) / 1e6) def test_openai_long_context_applies_to_the_whole_request(self): result = self.cost('openai', 'gpt-6-luna', {'input_tokens': 300_000, 'output_tokens': 1000}) self.assertAlmostEqual(result['total_cost'], (300_000 * 0.20 + 1000 * 0.75) / 1e6) def test_groq_cache_discount(self): result = self.cost('groq', 'openai/gpt-oss-120b', {'prompt_tokens': 1000, 'completion_tokens': 100, 'prompt_tokens_details': {'cached_tokens': 600}}) self.assertAlmostEqual(result['total_cost'], (400 * 0.15 + 600 * 0.075 + 100 * 0.60) / 1e6) def test_cached_tokens_of_a_model_without_cache_price_cost_the_input_price(self): result = self.cost('mistral', 'mistral-medium-latest', {'prompt_tokens': 1000, 'completion_tokens': 0, 'prompt_tokens_details': {'cached_tokens': 500}}) self.assertAlmostEqual(result['total_cost'], 1000 * 1.50 / 1e6) self.assertEqual(result['cache_savings'], 0.0) def test_announced_price_change_applies_from_its_day(self): before = catalog.price_for('google', 'gemini-3.8-flash', datetime.date(2026, 12, 31)) after = catalog.price_for('google', 'gemini-3.8-flash', datetime.date(2027, 1, 1)) self.assertEqual((before['input'], before['cached_input'], before['output']), (0.75, 0.075, 3.75)) self.assertEqual((after['input'], after['cached_input'], after['output']), (1.50, 0.15, 7.50)) self.assertNotIn('from', after) def test_unknown_model_keeps_tokens_and_costs_nothing(self): result = self.cost('openai', 'no-such-model', {'prompt_tokens': 10, 'completion_tokens': 5}) self.assertEqual((result['total_cost'], result['input_tokens'], result['output_tokens']), (0.0, 10, 5)) def test_historical_model_is_still_priced(self): self.assertGreater(self.cost('openai', 'gpt-4o', {'prompt_tokens': 1000, 'completion_tokens': 10})['total_cost'], 0) def test_czech_crowns_use_the_rate(self): tracker._rate_memo.update(day=datetime.date.today().isoformat(), rate=20.0) self.assertEqual(tracker.format_cost(1.0), '$1.00 (20 Kč)') class CatalogTests(unittest.TestCase): def test_every_offered_model_has_a_complete_price(self): for provider, models in catalog.MODELS.items(): for model in models: price = catalog.price_for(provider, model) self.assertTrue(price and price['input'] > 0 and price['output'] > 0, (provider, model)) for key in ('cached_input', 'cache_write'): if key in price: self.assertLessEqual(price['cached_input'], price['input']) self.assertEqual(set(catalog.PRICES), set(catalog.MODELS)) def test_prices_exist_only_for_offered_models(self): for provider, prices in catalog.PRICES.items(): self.assertEqual(set(prices), set(catalog.MODELS[provider]), provider) def test_settings_lists_and_defaults_match_the_catalog(self): with open(SETTINGS_XML, encoding='utf-8') as handle: text = handle.read() for provider, setting_id in catalog.SETTING_ID.items(): line = re.search(r']*id="%s"[^>]*>' % setting_id, text).group(0) self.assertEqual(re.search(r'values="([^"]*)"', line).group(1).split('|'), catalog.MODELS[provider]) self.assertEqual(re.search(r'default="([^"]*)"', line).group(1), catalog.DEFAULT_MODEL[provider]) if __name__ == '__main__': unittest.main()