|
| 1 | +//! ai_cost.rs 单测（AAA 模式；单价 env 读取为宿主环境——测试隔离保存/还原）。 |
| 2 | +//! @ai-context: 两个测试读写同一 env 变量——并行会竞态，用模块级互斥串行化。 |
| 3 | +
|
| 4 | +use std::sync::Mutex; |
| 5 | + |
| 6 | +use crate::ai_cost::{estimate_cost, estimate_for_content, estimate_tokens, price_per_1m, DEFAULT_PRICE_PER_1M}; |
| 7 | + |
| 8 | +/// env 操作互斥（防并行测试互相覆盖 SILICONFLOW_PRICE_PER_1M_TOKENS）。 |
| 9 | +static ENV_LOCK: Mutex<()> = Mutex::new(()); |
| 10 | + |
| 11 | +/// 在锁内执行 env 敏感操作（保存/还原宿主环境）。 |
| 12 | +fn with_env_locked(f: impl FnOnce()) { |
| 13 | + let _guard = ENV_LOCK.lock().unwrap_or_else(|e| e.into_inner()); |
| 14 | + let saved = std::env::var("SILICONFLOW_PRICE_PER_1M_TOKENS").ok(); |
| 15 | + f(); |
| 16 | + match saved { |
| 17 | + Some(v) => std::env::set_var("SILICONFLOW_PRICE_PER_1M_TOKENS", v), |
| 18 | + None => std::env::remove_var("SILICONFLOW_PRICE_PER_1M_TOKENS"), |
| 19 | + } |
| 20 | +} |
| 21 | + |
| 22 | +#[test] |
| 23 | +fn token_estimate_is_char_count() { |
| 24 | + // 中文 1 字符 ≈ 1 token（保守上界） |
| 25 | + assert_eq!(estimate_tokens(0), 0); |
| 26 | + assert_eq!(estimate_tokens(100), 100); |
| 27 | + assert_eq!(estimate_tokens(8000), 8000); |
| 28 | +} |
| 29 | + |
| 30 | +#[test] |
| 31 | +fn cost_zero_when_free_tier_or_bad_price() { |
| 32 | + // 免费档 / 非法单价（负数/NaN/inf）→ 0 费用（不产生负成本） |
| 33 | + assert_eq!(estimate_cost(1000, 0.0), 0.0); |
| 34 | + assert_eq!(estimate_cost(1000, -5.0), 0.0); |
| 35 | + assert_eq!(estimate_cost(1000, f64::NAN), 0.0); |
| 36 | + assert_eq!(estimate_cost(1000, f64::INFINITY), 0.0); |
| 37 | +} |
| 38 | + |
| 39 | +#[test] |
| 40 | +fn cost_scales_with_tokens_and_price() { |
| 41 | + // 1M token @ ¥10 → 每 1000 token = ¥0.01 |
| 42 | + let cost = estimate_cost(1000, 10.0); |
| 43 | + assert!((cost - 0.01).abs() < 1e-9, "实得 {}", cost); |
| 44 | + let cost2 = estimate_cost(500_000, 10.0); |
| 45 | + assert!((cost2 - 5.0).abs() < 1e-9); |
| 46 | +} |
| 47 | + |
| 48 | +#[test] |
| 49 | +fn price_default_free_tier() { |
| 50 | + // 缺省单价 = 免费档 0（当前默认模型免费） |
| 51 | + assert_eq!(DEFAULT_PRICE_PER_1M, 0.0); |
| 52 | + with_env_locked(|| { |
| 53 | + // env 未设置时 price_per_1m 回退默认 |
| 54 | + std::env::remove_var("SILICONFLOW_PRICE_PER_1M_TOKENS"); |
| 55 | + assert_eq!(price_per_1m(), 0.0); |
| 56 | + // 非法值回退默认 |
| 57 | + std::env::set_var("SILICONFLOW_PRICE_PER_1M_TOKENS", "abc"); |
| 58 | + assert_eq!(price_per_1m(), 0.0); |
| 59 | + // 合法覆盖生效 |
| 60 | + std::env::set_var("SILICONFLOW_PRICE_PER_1M_TOKENS", "12.5"); |
| 61 | + assert!((price_per_1m() - 12.5).abs() < 1e-9); |
| 62 | + }); |
| 63 | +} |
| 64 | + |
| 65 | +#[test] |
| 66 | +fn estimate_for_content_composes() { |
| 67 | + with_env_locked(|| { |
| 68 | + std::env::set_var("SILICONFLOW_PRICE_PER_1M_TOKENS", "10"); |
| 69 | + let est = estimate_for_content(5000); |
| 70 | + assert_eq!(est.est_tokens, 5000); |
| 71 | + assert!((est.est_cost_yuan - 0.05).abs() < 1e-9); |
| 72 | + assert_eq!(est.price_per_1m, 10.0); |
| 73 | + }); |
| 74 | +} |
0 commit comments