ext/goal/tests/accounting.rs
68 lines
#![allow(dead_code)]#[path = "../src/accounting.rs"]mod accounting;use accounting::GoalAccountingState;use codex_protocol::config_types::ModeKind;use codex_protocol::protocol::TokenUsage;use pretty_assertions::assert_eq;#[test]fn goal_accounting_uses_turn_start_baseline_for_exact_deltas() { let state = GoalAccountingState::default(); state.start_turn( "turn-1", ModeKind::Default, &token_usage( /*input_tokens*/ 100, /*cached_input_tokens*/ 10, /*output_tokens*/ 30, /*reasoning_output_tokens*/ 5, /*total_tokens*/ 135, ), ); let recorded = state .record_token_usage( "turn-1", &token_usage( /*input_tokens*/ 120, /*cached_input_tokens*/ 14, /*output_tokens*/ 42, /*reasoning_output_tokens*/ 8, /*total_tokens*/ 162, ), ) .expect("token delta should be recorded"); assert_eq!(28, recorded.turn_delta); assert_eq!(28, recorded.thread_unflushed_delta);}#[test]fn goal_accounting_ignores_plan_mode_turns() { let state = GoalAccountingState::default(); state.start_turn("turn-1", ModeKind::Plan, &TokenUsage::default()); let recorded = state.record_token_usage( "turn-1", &token_usage( /*input_tokens*/ 20, /*cached_input_tokens*/ 5, /*output_tokens*/ 8, /*reasoning_output_tokens*/ 2, /*total_tokens*/ 30, ), ); assert_eq!(None, recorded);}fn token_usage( input_tokens: i64, cached_input_tokens: i64, output_tokens: i64, reasoning_output_tokens: i64, total_tokens: i64,) -> TokenUsage { TokenUsage { input_tokens, cached_input_tokens, output_tokens, reasoning_output_tokens, total_tokens, }}