ed
ada9617308
Batch rename of 22 test files. 62 references renamed total. The full test suite is now GREEN again, matching the pre-rename baseline from Task 1.1. Pure mechanical rename. No behavior change. Files affected: test_ai_cache_tracking, test_ai_client_cli, test_ai_client_result, test_api_events, test_context_pruner, test_deepseek_provider, test_gemini_cli_* (3 files), test_gui2_mcp, test_headless_* (2 files), test_live_gui_integration_v2, test_orchestration_logic, test_phase6_engine, test_rag_integration, test_run_worker_lifecycle_abort, test_spawn_interception_v2, test_symbol_parsing, test_tier4_interceptor, test_tiered_aggregation, test_token_usage. Note: spec estimated 24 files; actual is 22 (test_deprecation_warnings no longer exists, and 1 fewer file than spec's list). Refs: conductor/tracks/send_result_to_send_20260616/
44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
import sys
|
|
import os
|
|
from unittest.mock import patch, MagicMock
|
|
from types import SimpleNamespace
|
|
|
|
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
|
|
|
|
from src import ai_client
|
|
from src.result_types import Result
|
|
|
|
def test_token_usage_tracking() -> None:
|
|
ai_client.reset_session()
|
|
ai_client.clear_comms_log()
|
|
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")
|
|
result = ai_client.send("Context", "Hello")
|
|
assert result.ok
|
|
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
|