use std::collections::BTreeMap;
use std::fs::File;
use std::io::BufReader;
use std::path::{Path, PathBuf};
use llmp::{Config, Record, Row, analyze, read_jsonl};
use serde_json::Value;
fn fixtures() -> Vec<PathBuf> {
let root = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures");
let mut files = Vec::new();
for dir in std::fs::read_dir(&root).unwrap() {
let dir = dir.unwrap().path();
if !dir.is_dir() {
continue;
}
for file in std::fs::read_dir(&dir).unwrap() {
let file = file.unwrap().path();
if file.extension().is_some_and(|e| e == "jsonl") {
files.push(file);
}
}
}
files.sort();
assert!(files.len() >= 4, "fixtures missing: {files:?}");
files
}
fn load(path: &Path) -> Vec<Record> {
let read = read_jsonl(BufReader::new(File::open(path).unwrap()));
assert!(
read.bad_lines.is_empty(),
"{}: {:?}",
path.display(),
read.bad_lines
);
read.records
}
fn golden_path(path: &Path) -> PathBuf {
let name = path
.file_name()
.unwrap()
.to_str()
.unwrap()
.trim_end_matches(".jsonl");
path.with_file_name(format!("{name}.report.json"))
}
#[test]
fn reports_match_the_golden_files() {
let update = std::env::var_os("LLMP_UPDATE_GOLDEN").is_some();
let mut stale = Vec::new();
for path in fixtures() {
let report = analyze(&load(&path), &Config::default());
let actual = serde_json::to_string_pretty(&report).unwrap() + "\n";
let golden = golden_path(&path);
if update {
std::fs::write(&golden, &actual).unwrap();
continue;
}
let expected = std::fs::read_to_string(&golden).unwrap_or_default();
if expected != actual {
stale.push(golden.display().to_string());
}
}
assert!(
stale.is_empty(),
"reports differ from {stale:?}; if the change is intended, rerun with LLMP_UPDATE_GOLDEN=1 and explain it in the commit"
);
}
#[test]
fn no_row_loses_more_than_it_evaluated() {
for path in fixtures() {
let report = analyze(&load(&path), &Config::default());
for row in &report.rows {
let evaluated = row.counts.evaluated_tokens.map_or(0, |e| e.value);
let lost: u32 = [&row.lost_client, &row.lost_external]
.into_iter()
.flatten()
.map(|l| l.tokens)
.sum();
assert!(
lost <= evaluated,
"{} row {}: lost {lost} > evaluated {evaluated}",
path.display(),
row.id
);
}
}
}
#[test]
fn every_number_is_finite() {
fn walk(value: &Value, at: &str) {
match value {
Value::Number(n) => assert!(n.as_f64().is_some_and(f64::is_finite), "{at}: {n}"),
Value::Array(items) => items
.iter()
.enumerate()
.for_each(|(i, v)| walk(v, &format!("{at}[{i}]"))),
Value::Object(map) => map.iter().for_each(|(k, v)| walk(v, &format!("{at}.{k}"))),
_ => {}
}
}
for path in fixtures() {
let report = analyze(&load(&path), &Config::default());
let text = serde_json::to_string(&report).unwrap();
assert!(
!text.contains("NaN") && !text.contains("-0.0"),
"{}",
path.display()
);
walk(
&serde_json::to_value(&report).unwrap(),
&path.display().to_string(),
);
}
}
#[test]
fn a_measured_loss_has_every_counter_and_an_unmeasured_one_has_no_loss() {
for path in fixtures() {
let report = analyze(&load(&path), &Config::default());
for row in &report.rows {
let c = &row.counts;
if row.loss_measured {
assert!(
c.prompt_tokens.is_some()
&& c.reused_tokens.is_some()
&& c.evaluated_tokens.is_some()
&& c.prefill_ms.is_some(),
"{} row {}",
path.display(),
row.id
);
} else {
assert!(
row.lost_client.is_none() && row.lost_external.is_none(),
"{} row {}",
path.display(),
row.id
);
}
}
}
}
fn request_key(record: &Record) -> Option<String> {
let mut body = record.request.body.clone()?;
let object = body.as_object_mut()?;
object.remove("stream");
object.remove("stream_options");
Some(body.to_string())
}
#[test]
fn a_stream_merges_to_the_same_totals_as_the_non_streamed_response() {
let mut pairs = 0;
for path in fixtures()
.into_iter()
.filter(|p| !p.to_str().unwrap().contains(".digest."))
{
let records = load(&path);
let report = analyze(&records, &Config::default());
let rows: BTreeMap<u64, &Row> = report.rows.iter().map(|r| (r.id, r)).collect();
let mut by_request: BTreeMap<String, Vec<&Record>> = BTreeMap::new();
for record in &records {
if let Some(key) = request_key(record) {
by_request.entry(key).or_default().push(record);
}
}
for group in by_request.values() {
let streamed = group
.iter()
.find(|r| r.response.as_ref().is_some_and(|x| x.stream.is_some()));
let plain = group
.iter()
.find(|r| r.response.as_ref().is_some_and(|x| x.stream.is_none()));
let (Some(streamed), Some(plain)) = (streamed, plain) else {
continue;
};
pairs += 1;
let (s, p) = (&rows[&streamed.id].counts, &rows[&plain.id].counts);
for (name, a, b) in [
("prompt_tokens", s.prompt_tokens, p.prompt_tokens),
(
"completion_tokens",
s.completion_tokens,
p.completion_tokens,
),
] {
if let (Some(a), Some(b)) = (a, b) {
assert_eq!(
a.value,
b.value,
"{}: {name} of records {} and {}",
path.display(),
streamed.id,
plain.id
);
}
}
let merged = streamed
.response
.as_ref()
.and_then(|r| r.body.as_ref())
.unwrap();
let reply = plain
.response
.as_ref()
.and_then(|r| r.body.as_ref())
.unwrap();
let shape = |v: &Value| {
let mut keys: Vec<String> = v.as_object().unwrap().keys().cloned().collect();
keys.retain(|k| {
!matches!(
k.as_str(),
"usage" | "timings" | "created_at" | "id" | "created"
)
});
keys.sort();
keys
};
let merged_shape = shape(merged);
let reply_shape = shape(reply);
assert!(
reply_shape
.iter()
.all(|k| merged_shape.contains(k) || k == "system_fingerprint"),
"{}: merged stream {merged_shape:?} lacks keys of {reply_shape:?}",
path.display()
);
}
}
assert!(
pairs >= 3,
"each full fixture should repeat one request streamed, found {pairs} pairs"
);
}
#[test]
fn full_and_digest_recordings_of_the_same_traffic_agree() {
let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/llama.cpp");
let full = analyze(&load(&dir.join("session.jsonl")), &Config::default());
let digest = analyze(&load(&dir.join("session.digest.jsonl")), &Config::default());
assert_eq!(full.rows.len(), digest.rows.len());
assert_eq!(full.coverage.analysed, digest.coverage.analysed);
assert_eq!(full.tools, digest.tools);
for (a, b) in full.rows.iter().zip(&digest.rows) {
assert_eq!(
(a.id, a.kind, a.session, a.turn),
(b.id, b.kind, b.session, b.turn)
);
assert_eq!(
(a.counts.prompt_tokens, a.counts.reused_tokens),
(b.counts.prompt_tokens, b.counts.reused_tokens)
);
assert_eq!(
(a.prefix_tokens, a.break_cause, a.cause),
(b.prefix_tokens, b.break_cause, b.cause),
"row {}",
a.id
);
assert_eq!(a.composition, b.composition, "row {}", a.id);
assert_eq!(a.lost_external, b.lost_external, "row {}", a.id);
assert_eq!(
a.lost_client.is_some(),
b.lost_client.is_some(),
"row {}",
a.id
);
}
}
#[test]
fn parallel_agents_are_linked_to_their_own_conversations() {
let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/llama.cpp");
for name in ["parallel.jsonl", "parallel-nocache.jsonl"] {
let records = load(&dir.join(name));
let report = analyze(&records, &Config::default());
let agent_of: BTreeMap<u64, (String, String)> = records
.iter()
.filter_map(|r| {
let body = r.request.body.as_ref()?;
let system = body["messages"][0]["content"]
.as_str()?
.chars()
.take(60)
.collect();
Some((r.id, (body["model"].as_str()?.to_string(), system)))
})
.collect();
let mut conversations: BTreeMap<(String, String), BTreeMap<u32, usize>> = BTreeMap::new();
for row in report.rows.iter().filter(|r| r.conversation.is_some()) {
*conversations
.entry(agent_of[&row.id].clone())
.or_default()
.entry(row.conversation.unwrap())
.or_default() += 1;
}
assert_eq!(conversations.len(), 5, "{name}: {conversations:?}");
for (agent, ids) in &conversations {
assert_eq!(ids.len(), 1, "{name}: {agent:?} split into {ids:?}");
}
assert!(
report.rows.iter().all(|r| r.break_cause.is_none()
|| r.break_cause == Some(llmp::BreakCause::SystemChanged)),
"{name}: interleaving produced a false break"
);
let clients: Vec<&str> = report.clients.iter().map(|c| c.client.as_str()).collect();
assert_eq!(
clients,
vec![
"llmp-drive/1 #1",
"llmp-drive/1 #2",
"other-agent/2",
"llmp-drive/1"
],
"{name}"
);
}
}