use std::{collections::BTreeMap, thread};
use kcode_k1_accounting::{Accounting, AccountingEvent, UsageValue};
use rust_decimal::Decimal;
fn event(source: &str, operation: &str, values: &[(&str, Decimal, Decimal)]) -> AccountingEvent {
let usage = values
.iter()
.map(|(key, units, price_cents)| {
(
(*key).to_owned(),
UsageValue {
units: *units,
price_cents: *price_cents,
},
)
})
.collect::<BTreeMap<_, _>>();
AccountingEvent {
source: source.to_owned(),
operation: operation.to_owned(),
usage,
}
}
#[test]
fn chronology_is_preserved() {
let accounting = Accounting::new();
let first = event("first", "create", &[]);
let second = event("second", "read", &[]);
accounting.record(&first);
accounting.record(&second);
assert_eq!(accounting.entries(), vec![first, second]);
}
#[test]
fn totals_aggregate_exact_case_sensitive_keys() {
let accounting = Accounting::new();
accounting.record(&event(
"provider-a",
"request",
&[
("input_tokens", Decimal::new(125, 2), Decimal::new(15, 1)),
("Input_Tokens", Decimal::new(2, 0), Decimal::new(3, 0)),
],
));
accounting.record(&event(
"provider-b",
"response",
&[("input_tokens", Decimal::new(25, 2), Decimal::new(5, 1))],
));
let totals = accounting.totals();
assert_eq!(
totals["input_tokens"],
UsageValue {
units: Decimal::new(15, 1),
price_cents: Decimal::new(2, 0),
}
);
assert_eq!(
totals["Input_Tokens"],
UsageValue {
units: Decimal::new(2, 0),
price_cents: Decimal::new(3, 0),
}
);
}
#[test]
fn decimal_values_remain_exact() {
let accounting = Accounting::new();
accounting.record(&event(
"provider",
"request",
&[("requests", Decimal::new(1, 1), Decimal::new(1, 1))],
));
accounting.record(&event(
"provider",
"request",
&[("requests", Decimal::new(2, 1), Decimal::new(2, 1))],
));
assert_eq!(
accounting.totals()["requests"],
UsageValue {
units: Decimal::new(3, 1),
price_cents: Decimal::new(3, 1),
}
);
}
#[test]
fn empty_usage_is_recorded() {
let accounting = Accounting::new();
let attempted = event("provider", "attempt", &[]);
accounting.record(&attempted);
assert_eq!(accounting.entries(), vec![attempted]);
assert!(accounting.totals().is_empty());
}
#[test]
fn cloned_handles_share_a_ledger() {
let accounting = Accounting::new();
let clone = accounting.clone();
clone.record(&event(
"provider",
"request",
&[("requests", Decimal::ONE, Decimal::ONE)],
));
assert_eq!(accounting.entries().len(), 1);
}
#[test]
fn concurrent_recording() {
let accounting = Accounting::new();
thread::scope(|scope| {
for worker in 0..8 {
let handle = accounting.clone();
scope.spawn(move || {
let recorded = event(
&format!("worker-{worker}"),
"request",
&[("requests", Decimal::ONE, Decimal::ONE)],
);
for _ in 0..100 {
handle.record(&recorded);
}
});
}
});
assert_eq!(accounting.entries().len(), 800);
assert_eq!(
accounting.totals()["requests"],
UsageValue {
units: Decimal::new(800, 0),
price_cents: Decimal::new(800, 0),
}
);
}