use super::{Fullness, of_agent};
use crate::budgets::{BudgetSpend, StepBill};
use std::collections::BTreeMap;
const ROOT: &str = "20260802T120000Z-root";
const KID: &str = "20260802T120000Z-root-20260802T120100Z-kid0";
fn usage(input: u64, read: u64, write: u64) -> BudgetSpend {
BudgetSpend {
input_tokens: input,
output_tokens: 7,
cache_read_tokens: read,
cache_write_tokens: write,
}
}
fn bill(conv: &str, seq: &str, model: Option<&str>, last: BudgetSpend) -> StepBill {
StepBill {
conv: conv.to_owned(),
seq: seq.to_owned(),
model: model.map(str::to_owned),
spend: usage(999, 999, 999),
last_usage: last,
wall_secs: 0,
}
}
fn windows() -> BTreeMap<String, u64> {
BTreeMap::from([("sonnet".to_owned(), 200_000)])
}
#[test]
fn fullness_reads_the_folded_prompt_not_the_bare_input_counter() {
let bills = [bill(ROOT, "001", Some("sonnet"), usage(2_000, 48_000, 0))];
let full = of_agent(&bills, ROOT, &windows()).expect("a measured context");
assert_eq!(full.prompt_tokens, 48_000);
}
#[test]
fn the_root_s_latest_step_is_the_context() {
let bills = [
bill(ROOT, "001", Some("sonnet"), usage(10_000, 0, 0)),
bill(ROOT, "002", Some("sonnet"), usage(2_000, 48_000, 0)),
bill(KID, "009", Some("sonnet"), usage(190_000, 0, 0)),
];
let full = of_agent(&bills, ROOT, &windows()).expect("a measured context");
assert_eq!(
full,
Fullness {
model: "sonnet".to_owned(),
prompt_tokens: 48_000,
window: 200_000,
}
);
assert_eq!(full.percent(), 24);
}
#[test]
fn step_order_is_lexical_because_the_width_is_fixed() {
let bills = [
bill(ROOT, "009", Some("sonnet"), usage(9, 0, 0)),
bill(ROOT, "010", Some("sonnet"), usage(100_000, 0, 0)),
];
let full = of_agent(&bills, ROOT, &windows()).expect("a measured context");
assert_eq!(full.prompt_tokens, 100_000);
assert_eq!(full.percent(), 50);
}
#[test]
fn nothing_honest_to_say_is_no_figure_at_all() {
assert_eq!(of_agent(&[], ROOT, &windows()), None);
assert_eq!(
of_agent(
&[bill(ROOT, "001", None, usage(10, 0, 0))],
ROOT,
&windows()
),
None
);
assert_eq!(
of_agent(
&[bill(ROOT, "001", Some("opus"), usage(10, 0, 0))],
ROOT,
&windows()
),
None
);
}
#[test]
fn a_context_past_its_declared_window_is_not_clamped() {
let bills = [bill(ROOT, "001", Some("sonnet"), usage(280_000, 0, 0))];
let full = of_agent(&bills, ROOT, &windows()).expect("a measured context");
assert_eq!(full.percent(), 140);
}