# -*- coding: utf-8 -*- """Unit tests for the token usage core module.""" from __future__ import annotations import asyncio import json from unittest.mock import MagicMock import pytest from pineagents.token_usage.buffer import ( TokenUsageBuffer, _UsageEvent, _apply_event, ) from pineagents.token_usage.manager import ( TokenUsageByDateModel, TokenUsageByModel, TokenUsageManager, TokenUsageRecord, TokenUsageStats, TokenUsageSummary, ) from pineagents.token_usage.model_wrapper import TokenRecordingModelWrapper from pineagents.token_usage.storage import load_data, save_data_sync # ============================================================================= # Fixtures # ============================================================================= @pytest.fixture(autouse=True) def _isolate_token_usage_manager(): """Isolate token usage manager singleton for each test.""" # pylint: disable=protected-access TokenUsageManager._instance = None yield TokenUsageManager._instance = None # ============================================================================= # Test _apply_event # ============================================================================= class TestApplyEvent: """Test the _apply_event function that accumulates usage events.""" # pylint: disable=protected-access def test_apply_event_creates_new_entry(self): """Should create new entry for first event.""" cache = {} event = _UsageEvent( provider_id="openai", model_name="gpt-4", prompt_tokens=100, completion_tokens=50, date_str="2026-04-24", now_iso="2026-04-24T10:00:00+00:00", ) _apply_event(cache, event) assert "2026-04-24" in cache assert "openai:gpt-4" in cache["2026-04-24"] entry = cache["2026-04-24"]["openai:gpt-4"] assert entry["prompt_tokens"] == 100 assert entry["completion_tokens"] == 50 assert entry["call_count"] == 1 def test_apply_event_accumulates_same_model(self): """Should accumulate tokens for same provider:model on same date.""" cache = {} for _ in range(3): _apply_event( cache, _UsageEvent( provider_id="openai", model_name="gpt-4", prompt_tokens=100, completion_tokens=50, date_str="2026-04-24", now_iso="2026-04-24T10:00:00+00:00", ), ) entry = cache["2026-04-24"]["openai:gpt-4"] assert entry["prompt_tokens"] == 300 assert entry["call_count"] == 3 # ============================================================================= # Test Storage # ============================================================================= class TestStorage: """Test storage load/save operations.""" @pytest.mark.asyncio async def test_load_data_nonexistent_file(self, tmp_path): """Should return empty dict when file doesn't exist.""" data = await load_data(tmp_path / "token_usage.json") assert data == {} @pytest.mark.asyncio async def test_load_data_valid_json(self, tmp_path): """Should load and return valid JSON data.""" path = tmp_path / "token_usage.json" expected = { "2026-04-24": { "openai:gpt-4": { "provider_id": "openai", "model_name": "gpt-4", "prompt_tokens": 100, "completion_tokens": 50, "call_count": 2, }, }, } path.write_text(json.dumps(expected)) data = await load_data(path) assert data["2026-04-24"]["openai:gpt-4"]["prompt_tokens"] == 100 @pytest.mark.asyncio async def test_load_data_corrupt_json(self, tmp_path): """Should handle corrupt JSON gracefully.""" path = tmp_path / "token_usage.json" path.write_text("{invalid json}") data = await load_data(path) assert data == {} def test_save_data_sync_writes_file(self, tmp_path): """Should write data to file atomically.""" path = tmp_path / "token_usage.json" data = {"2026-04-24": {"openai:gpt-4": {"prompt_tokens": 100}}} save_data_sync(path, data) assert path.exists() loaded = json.loads(path.read_text()) assert loaded == data def test_save_data_sync_creates_parent_dirs(self, tmp_path): """Should create parent directories if needed.""" path = tmp_path / "subdir" / "token_usage.json" save_data_sync(path, {"test": "data"}) assert path.exists() # ============================================================================= # Test TokenUsageBuffer # ============================================================================= class TestTokenUsageBuffer: """Test TokenUsageBuffer core functionality.""" # pylint: disable=protected-access def test_init_defaults(self, tmp_path): """Should initialize with correct defaults.""" buffer = TokenUsageBuffer(tmp_path / "test.json") assert buffer._flush_interval == 10 @pytest.mark.asyncio async def test_enqueue_adds_to_queue(self, tmp_path): """Should add event to queue.""" buffer = TokenUsageBuffer(tmp_path / "test.json") event = _UsageEvent( provider_id="openai", model_name="gpt-4", prompt_tokens=100, completion_tokens=50, date_str="2026-04-24", now_iso="2026-04-24T10:00:00+00:00", ) buffer.enqueue(event) assert buffer._queue.qsize() == 1 @pytest.mark.asyncio async def test_consumer_processes_events(self, tmp_path): """Consumer should process and accumulate events.""" buffer = TokenUsageBuffer(tmp_path / "test.json") buffer.start() for _ in range(3): buffer.enqueue( _UsageEvent( provider_id="openai", model_name="gpt-4", prompt_tokens=100, completion_tokens=50, date_str="2026-04-24", now_iso="2026-04-24T10:00:00+00:00", ), ) await asyncio.sleep(0.2) await buffer.stop() entry = buffer._disk_cache["2026-04-24"]["openai:gpt-4"] assert entry["prompt_tokens"] == 300 assert entry["call_count"] == 3 @pytest.mark.asyncio async def test_stop_does_not_wipe_history_when_seed_interrupted( self, tmp_path, monkeypatch, ): """A stop() that races cache seeding must not clobber the file.""" path = tmp_path / "test.json" existing = { "2026-04-24": { "openai:gpt-4": { "provider_id": "openai", "model_name": "gpt-4", "prompt_tokens": 100, "completion_tokens": 50, "call_count": 1, }, }, } path.write_text(json.dumps(existing), encoding="utf-8") seeding = asyncio.Event() async def _never_returns(_path): # Park the consumer inside the seed so stop() runs while # ``_disk_cache`` is still the initial empty dict. seeding.set() await asyncio.Event().wait() return {} monkeypatch.setattr( "pineagents.token_usage.buffer.load_data", _never_returns, ) buffer = TokenUsageBuffer(path, flush_interval=3600) buffer.start() await asyncio.wait_for(seeding.wait(), timeout=1) await buffer.stop() assert json.loads(path.read_text(encoding="utf-8")) == existing @pytest.mark.asyncio async def test_stop_flushes_after_seed_completes(self, tmp_path): """Normal shutdown still merges new events into stored history.""" path = tmp_path / "test.json" path.write_text( json.dumps( { "2026-04-23": { "openai:gpt-4": { "provider_id": "openai", "model_name": "gpt-4", "prompt_tokens": 7, "completion_tokens": 3, "call_count": 1, }, }, }, ), encoding="utf-8", ) buffer = TokenUsageBuffer(path, flush_interval=3600) buffer.start() buffer.enqueue( _UsageEvent( provider_id="openai", model_name="gpt-4", prompt_tokens=100, completion_tokens=50, date_str="2026-04-24", now_iso="2026-04-24T10:00:00+00:00", ), ) await buffer.stop() written = json.loads(path.read_text(encoding="utf-8")) assert written["2026-04-23"]["openai:gpt-4"]["prompt_tokens"] == 7 assert written["2026-04-24"]["openai:gpt-4"]["prompt_tokens"] == 100 # ============================================================================= # Test Pydantic Models # ============================================================================= class TestTokenUsageStats: """Test TokenUsageStats model.""" def test_default_values(self): """Should have zero defaults.""" stats = TokenUsageStats() assert stats.prompt_tokens == 0 assert stats.completion_tokens == 0 assert stats.call_count == 0 def test_custom_values(self): """Should accept custom values.""" stats = TokenUsageStats( prompt_tokens=100, completion_tokens=50, call_count=5, ) assert stats.prompt_tokens == 100 assert stats.completion_tokens == 50 assert stats.call_count == 5 def test_validation_rejects_negative(self): """Should reject negative values.""" with pytest.raises(Exception): TokenUsageStats(prompt_tokens=-1) class TestTokenUsageModels: """Test TokenUsage models.""" def test_create_record(self): """Should create record with all fields.""" record = TokenUsageRecord( date="2026-04-24", provider_id="openai", model="gpt-4", prompt_tokens=100, completion_tokens=50, call_count=3, ) assert record.date == "2026-04-24" assert record.provider_id == "openai" assert record.model == "gpt-4" def test_empty_summary(self): """Should create empty summary with defaults.""" summary = TokenUsageSummary() assert summary.total_prompt_tokens == 0 assert summary.total_completion_tokens == 0 assert summary.total_calls == 0 assert summary.by_model == {} assert summary.by_date == {} def test_summary_with_data(self): """Should accept populated data.""" summary = TokenUsageSummary( total_prompt_tokens=500, total_completion_tokens=250, total_calls=10, by_model={ "openai:gpt-4": TokenUsageByModel( provider_id="openai", model="gpt-4", prompt_tokens=500, completion_tokens=250, call_count=10, ), }, by_date={ "2026-04-24": TokenUsageStats( prompt_tokens=500, completion_tokens=250, call_count=10, ), }, ) assert summary.total_prompt_tokens == 500 assert len(summary.by_model) == 1 assert summary.by_model["openai:gpt-4"].model == "gpt-4" assert len(summary.by_date) == 1 def test_token_usage_by_model(self): """Should create TokenUsageByModel with provider_id.""" by_model = TokenUsageByModel( provider_id="openai", model="gpt-4", prompt_tokens=300, completion_tokens=150, call_count=6, ) assert by_model.provider_id == "openai" assert by_model.model == "gpt-4" def test_token_usage_by_date_model(self): """Should create TokenUsageByDateModel.""" by_date_model = TokenUsageByDateModel( provider_id="dashscope", model="qwen3-max", prompt_tokens=200, completion_tokens=100, call_count=4, ) assert by_date_model.provider_id == "dashscope" assert by_date_model.model == "qwen3-max" # ============================================================================= # Test TokenUsageManager # ============================================================================= class TestTokenUsageManagerCore: """Test TokenUsageManager singleton, lifecycle, and operations.""" def test_get_instance_returns_singleton(self, tmp_path, monkeypatch): """Should return same instance on multiple calls.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) manager1 = TokenUsageManager.get_instance() manager2 = TokenUsageManager.get_instance() assert manager1 is manager2 @pytest.mark.asyncio async def test_start_and_stop(self, tmp_path, monkeypatch): """Should start and stop cleanly.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) manager = TokenUsageManager() manager.start(flush_interval=10) await manager.stop() @pytest.mark.asyncio async def test_record_usage(self, tmp_path, monkeypatch): """Should record token usage.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) manager = TokenUsageManager() manager.start(flush_interval=10) await manager.record( provider_id="openai", model_name="gpt-4", prompt_tokens=100, completion_tokens=50, ) await asyncio.sleep(0.2) await manager.stop() @pytest.mark.asyncio async def test_get_summary_empty(self, tmp_path, monkeypatch): """Should return empty summary when no data.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) manager = TokenUsageManager() manager.start(flush_interval=10) summary = await manager.get_summary() assert summary.total_prompt_tokens == 0 assert summary.total_completion_tokens == 0 assert summary.total_calls == 0 assert summary.by_date == {} await manager.stop() @pytest.mark.asyncio async def test_get_details_empty(self, tmp_path, monkeypatch): """Should return empty list when no data.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) manager = TokenUsageManager() manager.start(flush_interval=10) details = await manager.get_details() assert details == [] await manager.stop() @pytest.mark.asyncio async def test_get_details_with_data(self, tmp_path, monkeypatch): """Should return raw records for frontend aggregation.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) manager = TokenUsageManager() manager.start(flush_interval=10) # Record some usage await manager.record( provider_id="openai", model_name="gpt-4", prompt_tokens=100, completion_tokens=50, ) await manager.record( provider_id="dashscope", model_name="qwen3-max", prompt_tokens=200, completion_tokens=100, ) await asyncio.sleep(0.2) details = await manager.get_details() # Should have 2 records assert len(details) == 2 # Verify structure models = {r.model for r in details} assert "gpt-4" in models assert "qwen3-max" in models await manager.stop() # ============================================================================= # Test TokenRecordingModelWrapper # ============================================================================= class TestTokenRecordingModelWrapper: """Test TokenRecordingModelWrapper.""" # pylint: disable=protected-access def test_init_wraps_model(self, tmp_path, monkeypatch): """Should wrap a ChatModelBase instance.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) mock_model = MagicMock() mock_model.model = "gpt-4" wrapper = TokenRecordingModelWrapper( provider_id="openai", model=mock_model, ) assert wrapper._provider_id == "openai" assert wrapper._model is mock_model assert wrapper.model == "gpt-4" def test_record_usage_with_valid_usage(self, tmp_path, monkeypatch): """Should record valid usage.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) mock_model = MagicMock() mock_model.model = "gpt-4" wrapper = TokenRecordingModelWrapper( provider_id="openai", model=mock_model, ) mock_usage = MagicMock() mock_usage.input_tokens = 100 mock_usage.output_tokens = 50 wrapper._record_usage(mock_usage) def test_record_usage_includes_context_and_threshold( self, tmp_path, monkeypatch, ): """Per-call usage carries context_size and compaction threshold.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", tmp_path, ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) monkeypatch.setattr( "pineagents.app.agent_context.get_current_session_id", lambda: "sess-1", ) mock_model = MagicMock() mock_model.model = "gpt-4" mock_model.context_size = 1_000_000 wrapper = TokenRecordingModelWrapper( provider_id="openai", model=mock_model, compact_threshold=0.8, ) mock_usage = MagicMock() mock_usage.input_tokens = 123_000 mock_usage.output_tokens = 50 wrapper._record_usage(mock_usage) stored = TokenRecordingModelWrapper.pop_usage_for_session("sess-1") assert stored is not None assert stored["context_size"] == 1_000_000 assert stored["compact_threshold"] == 0.8 def test_pop_usage_for_session(self, monkeypatch): """Should pop usage for session.""" monkeypatch.setattr( "pineagents.token_usage.manager.WORKING_DIR", "/tmp", ) monkeypatch.setattr( "pineagents.token_usage.manager.TOKEN_USAGE_FILE", "test_token_usage.json", ) # Clear any existing usage TokenRecordingModelWrapper._usage_by_session.clear() # Add test usage TokenRecordingModelWrapper._usage_by_session["test-session"] = { "prompt_tokens": 100, } usage = TokenRecordingModelWrapper.pop_usage_for_session( "test-session", ) assert usage is not None assert usage["prompt_tokens"] == 100 # Verify it was removed assert ( TokenRecordingModelWrapper.pop_usage_for_session("test-session") is None )