import sys import os import hashlib from unittest.mock import patch, MagicMock from types import SimpleNamespace from pathlib import Path sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))) from src import ai_client def test_token_usage_tracking() -> None: ai_client.reset_session() with patch("src.ai_client._ensure_gemini_client"), \ patch("src.ai_client._gemini_client") as mock_client: mock_chat = MagicMock() mock_client.chats.create.return_value = mock_chat mock_usage = SimpleNamespace( prompt_token_count=100, candidates_token_count=50, total_token_count=150, cached_content_token_count=20 ) mock_part = SimpleNamespace(text="Mock Response", function_call=None) mock_content = SimpleNamespace(parts=[mock_part]) mock_candidate = SimpleNamespace() mock_candidate.content = mock_content mock_candidate.finish_reason = SimpleNamespace(name="STOP") mock_response = SimpleNamespace() mock_response.candidates = [mock_candidate] mock_response.usage_metadata = mock_usage mock_response.text = "Mock Response" mock_chat.send_message.return_value = mock_response ai_client.set_provider("gemini", "gemini-2.5-flash-lite") ai_client.send("Context", "Hello") comms = ai_client.get_comms_log() response_entries = [e for e in comms if e.get("direction") == "IN" and e["kind"] == "response"] assert len(response_entries) > 0 usage = response_entries[0]["payload"]["usage"] assert usage["input_tokens"] == 100 assert usage["output_tokens"] == 50 assert usage["cache_read_input_tokens"] == 20