# -*- coding: utf-8 -*- """Unit tests: the AI usage history (wcs/ai/usage_store.py) and its recording (UsageTracker.track_ai_usage).""" import csv import datetime import json import os import sys import tempfile 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 usage_store as store # noqa: E402 D = datetime.datetime def record(**kwargs): values = dict(provider='openai', model='gpt-6-luna', purpose='millionaire_chat', input_tokens=100, cached_tokens=50, output_tokens=20, input_cost=0.01, cached_cost=0.002, output_cost=0.005, savings=0.001) values.update(kwargs) return store.record_call(**values) class StoreTests(unittest.TestCase): def setUp(self): self.dir = tempfile.mkdtemp(prefix='milionar-usage-') store.set_path(os.path.join(self.dir, 'ai_usage.db')) def test_totals_and_breakdowns(self): record(now=D(2026, 10, 1, 10)) record(now=D(2026, 10, 2, 10), model='gpt-6.1-sol') record(now=D(2026, 10, 2, 11), provider='anthropic', model='claude-opus-5-5', purpose='quiz') record(now=D(2026, 10, 2, 12), status='error', error='HTTP 500', input_tokens=0, cached_tokens=0, output_tokens=0, input_cost=0, cached_cost=0, output_cost=0, savings=0) total = store.totals() self.assertEqual((total['requests'], total['errors']), (4, 1)) self.assertAlmostEqual(total['cost'], 3 * 0.017) self.assertEqual(total['cached_tokens'], 150) self.assertEqual(store.totals(since='2026-10-02')['requests'], 3) self.assertEqual(store.totals(since='2026-10-02', until='2026-10-01')['requests'], 0) self.assertEqual({row['provider'] for row in store.per_provider()}, {'openai', 'anthropic'}) self.assertEqual({(row['provider'], row['model']) for row in store.per_model()}, {('openai', 'gpt-6-luna'), ('openai', 'gpt-6.1-sol'), ('anthropic', 'claude-opus-5-5')}) purposes = {row['purpose']: row['requests'] for row in store.per_purpose()} self.assertEqual(purposes, {'millionaire_chat': 3, 'quiz': 1}) self.assertEqual(store.span(), ('2026-10-01', '2026-10-02')) def test_per_day_has_every_day_of_the_range(self): record(now=D(2026, 10, 3, 9)) days = store.per_day('2026-10-01', '2026-10-05') self.assertEqual([d['day'] for d in days], ['2026-10-0%d' % n for n in range(1, 6)]) self.assertEqual([d['requests'] for d in days], [0, 0, 1, 0, 0]) def test_a_game_is_one_history_entry_and_a_loose_call_is_its_own(self): for minute in (0, 5, 9): record(now=D(2026, 10, 2, 20, minute), session_id='millionaire_abc', title='Matrix') record(now=D(2026, 10, 2, 21), purpose='quiz') entries = store.sessions() self.assertEqual(len(entries), 2) quiz, game = entries self.assertTrue(quiz['key'].startswith('c:')) self.assertEqual((game['key'], game['title'], game['requests']), ('s:millionaire_abc', 'Matrix', 3)) self.assertAlmostEqual(game['cost'], 3 * 0.017) self.assertEqual(len(store.calls_of('s:millionaire_abc')), 3) self.assertEqual(len(store.calls_of(quiz['key'])), 1) def test_deleting(self): for day in (1, 2): record(now=D(2026, 10, day, 10), session_id='g%d' % day) record(now=D(2026, 9, 30, 10), provider='google', model='gemini-3.8-flash') self.assertEqual(store.delete_entry('s:g1'), 1) self.assertEqual(store.totals()['requests'], 2) self.assertEqual(store.delete_day('2026-10-02'), 1) self.assertEqual(store.delete_provider('google'), 1) self.assertEqual(store.totals()['requests'], 0) record(now=D(2026, 10, 1, 10)) record(now=D(2026, 9, 1, 10)) self.assertEqual(store.delete_month('2026-09'), 1) self.assertEqual(store.delete_call(store.sessions()[0]['key'] and int(store.sessions()[0]['key'][2:])), 1) record() self.assertEqual(store.delete_all(), 1) self.assertEqual(store.totals()['requests'], 0) def test_a_delete_can_be_previewed_first(self): record(now=D(2026, 10, 1, 10), session_id='g1') record(now=D(2026, 10, 1, 11), session_id='g1') record(now=D(2026, 10, 2, 10), provider='groq', model='openai/gpt-oss-20b') record(now=D(2025, 1, 5, 10)) store.prune(datetime.date(2026, 10, 7)) self.assertEqual(store.preview_delete('entry', 's:g1')['requests'], 2) self.assertAlmostEqual(store.preview_delete('entry', 's:g1')['cost'], 2 * 0.017) self.assertEqual(store.preview_delete('month', '2025-01')['records'], 1) # a rolled-up day self.assertEqual(store.preview_delete('provider', 'groq')['records'], 1) self.assertEqual(store.preview_delete('all')['records'], 4) self.assertEqual(store.totals()['requests'], 4) # a preview deletes nothing with self.assertRaises(ValueError): store.preview_delete('week') def test_pruning_rolls_old_calls_up_without_changing_the_totals(self): record(now=D(2025, 8, 15, 10), session_id='old') record(now=D(2025, 8, 15, 11), session_id='old') record(now=D(2025, 8, 16, 11), provider='groq', model='openai/gpt-oss-20b', status='error') record(now=D(2026, 10, 1, 10)) before = store.totals() self.assertEqual(store.prune(datetime.date(2026, 10, 7)), 3) # the cutoff is 2025-10-01 after = store.totals() self.assertEqual((after['requests'], after['errors']), (before['requests'], before['errors'])) self.assertAlmostEqual(after['cost'], before['cost']) self.assertEqual(len(store.sessions()), 1) # only the recent call has its detail self.assertEqual(store.per_day('2025-08-15', '2025-08-15')[0]['requests'], 2) self.assertEqual(store.delete_month('2025-08'), 2) # rolled-up days can be deleted too self.assertEqual(store.totals()['requests'], 1) def test_pruning_runs_once_a_day(self): record(now=D(2025, 1, 1, 10)) self.assertEqual(store.prune_if_due(datetime.date(2026, 10, 7)), 1) record(now=D(2025, 1, 2, 10)) self.assertEqual(store.prune_if_due(datetime.date(2026, 10, 7)), 0) def test_export(self): record(now=D(2026, 10, 1, 10), title='Matrix, Reloaded') path = os.path.join(self.dir, 'out.csv') self.assertEqual(store.export_csv(path), 1) with open(path, encoding='utf-8', newline='') as handle: rows = list(csv.DictReader(handle)) self.assertEqual((rows[0]['title'], rows[0]['model'], rows[0]['requests']), ('Matrix, Reloaded', 'gpt-6-luna', '1')) class MigrationTests(unittest.TestCase): def test_the_old_json_becomes_rows_and_keeps_its_totals(self): directory = tempfile.mkdtemp(prefix='milionar-usage-') store.set_path(os.path.join(directory, 'ai_usage.db')) legacy = { 'created': '2026-01-05T10:00:00', 'providers': {'openai': {'models': {'gpt-6-luna': {'requests': 40, 'total_cost': 0.5, 'input_tokens': 4000, 'cached_tokens': 1000, 'output_tokens': 900}}}}, 'call_history': [{'timestamp': '2026-10-06T22:02:27', 'provider': 'openai', 'model': 'gpt-6-luna', 'request_type': 'millionaire_chat', 'session_id': 'millionaire_5e', 'inference_count': 4, 'cost': 0.1, 'input_tokens': 2000, 'cached_tokens': 500, 'output_tokens': 400}]} with open(os.path.join(directory, 'ai_usage_stats.json'), 'w', encoding='utf-8') as handle: json.dump(legacy, handle) total = store.totals() self.assertEqual(total['requests'], 40) self.assertAlmostEqual(total['cost'], 0.5) self.assertEqual((total['input_tokens'], total['cached_tokens'], total['output_tokens']), (4000, 1000, 900)) entry = store.sessions()[0] self.assertEqual((entry['key'], entry['requests'], entry['purpose']), ('s:millionaire_5e', 4, 'millionaire_chat')) self.assertTrue(os.path.exists(os.path.join(directory, 'ai_usage_stats.json.migrated'))) self.assertFalse(os.path.exists(os.path.join(directory, 'ai_usage_stats.json'))) self.assertEqual(store.totals()['requests'], 40) # not imported twice class TrackerTests(unittest.TestCase): def setUp(self): store.set_path(os.path.join(tempfile.mkdtemp(prefix='milionar-usage-'), 'ai_usage.db')) tracker._session_titles.clear() def test_a_call_is_recorded_with_its_price_latency_and_game_title(self): tracker.describe_session('millionaire_x', 'Matrix') tracker.track_ai_usage('openai', 'gpt-6-luna', {'input_tokens': 1000, 'output_tokens': 500, 'latency_ms': 1234, 'input_tokens_details': {'cached_tokens': 400}}, 'millionaire_chat', session_id='millionaire_x') call = store.calls_of('s:millionaire_x')[0] self.assertEqual((call['title'], call['latency_ms'], call['status'], call['purpose']), ('Matrix', 1234, 'ok', 'millionaire_chat')) self.assertAlmostEqual(call['cost'], (600 * 0.10 + 400 * 0.01 + 500 * 0.50) / 1e6) self.assertAlmostEqual(call['savings'], 400 * 0.09 / 1e6) self.assertEqual((call['input_tokens'], call['cached_tokens'], call['output_tokens']), (600, 400, 500)) def test_a_failed_call_is_recorded_with_its_reason(self): tracker.track_ai_failure('anthropic', 'claude-opus-5-5', 'quiz', RuntimeError('HTTP 529 overloaded'), latency_ms=800) call = store.sessions()[0] self.assertEqual((call['errors'], call['cost']), (1, 0)) self.assertEqual(store.calls_of(call['key'])[0]['error'], 'HTTP 529 overloaded') def test_missing_usage_still_records_the_call(self): tracker.track_ai_usage('mistral', 'mistral-medium-latest', None, 'quiz') self.assertEqual(store.totals()['requests'], 1) if __name__ == '__main__': unittest.main()