use std::error::Error as StdError;
use std::io::Cursor;
use std::num::NonZeroU32;
use haem_frames::embedding::{
Embed, Failed, FailedCode, Input, ItemCode, Kind, Message, Outcome, Tag, decode_message,
read_message, read_start, write_message,
};
use serde_json::{Value, json};
type TestResult = Result<(), Box<dyn StdError>>;
fn examples() -> Vec<(Tag, Value)> {
vec![
(
Tag::Start,
json!({"protocol":1,"model_dir":"/models/test","limits":{"memory_bytes":1024,"threads":2,"batch_items":1,"input_bytes":32,"tokens":16,"frame_bytes":8192}}),
),
(
Tag::Ready,
json!({"protocol":1,"worker":"worker/1","core_limit":0,"descriptors_closed":0,"environment":0,"model":{"name":"dense","revision":"r1","manifest_sha256":"a".repeat(64),"dimensions":1,"max_tokens":32,"special_tokens":2,"prefix_tokens":0,"normalised":true,"distance":"cosine"}}),
),
(
Tag::Embed,
json!({"kind":"document","items":[{"id":"one","text":"text"}]}),
),
(
Tag::Vectors,
json!({"items":[{"id":"one","outcome":"vector","vector":"0000803f","tokens_read":3,"tokens_total":3}]}),
),
(
Tag::Failed,
json!({"code":"embed_model_missing","message":"model is absent"}),
),
]
}
fn protocol(tag: Tag, value: &Value) -> TestResult {
let result = decode_message(tag as u8, &serde_json::to_vec(value)?);
assert_eq!(
result.err().ok_or("invalid message accepted")?.code,
FailedCode::EmbedProtocol
);
Ok(())
}
#[test]
fn every_message_round_trips_through_shared_framing() -> TestResult {
for (tag, value) in examples() {
let message = decode_message(tag as u8, &serde_json::to_vec(&value)?)?;
let mut bytes = Vec::new();
write_message(&mut bytes, &message, NonZeroU32::MAX)?;
let received =
read_message(&mut Cursor::new(bytes), NonZeroU32::MAX)?.ok_or("missing message")?;
assert_eq!(received.tag(), tag);
}
Ok(())
}
#[test]
fn ready_requires_prefix_tokens_and_accepts_zero() -> TestResult {
let mut value = examples().remove(1).1;
let Message::Ready(ready) = decode_message(Tag::Ready as u8, &serde_json::to_vec(&value)?)?
else {
return Err("not Ready".into());
};
assert_eq!(ready.model.prefix_tokens, 0);
assert!(
value["model"]
.as_object_mut()
.ok_or("model is not an object")?
.remove("prefix_tokens")
.is_some()
);
let error = decode_message(Tag::Ready as u8, &serde_json::to_vec(&value)?)
.err()
.ok_or("Ready without prefix_tokens was accepted")?;
assert_eq!(error.code, FailedCode::EmbedProtocol);
assert!(error.message.contains("prefix_tokens"));
Ok(())
}
#[test]
fn inference_failure_round_trips_by_its_exact_name() -> TestResult {
let value = json!({"code":"embed_inference_failed","message":"backend failed"});
let message = decode_message(Tag::Failed as u8, &serde_json::to_vec(&value)?)?;
let mut bytes = Vec::new();
write_message(&mut bytes, &message, NonZeroU32::MAX)?;
let Message::Failed(failed) =
read_message(&mut Cursor::new(bytes), NonZeroU32::MAX)?.ok_or("missing Failed")?
else {
return Err("not Failed".into());
};
assert_eq!(failed.code, FailedCode::EmbedInferenceFailed);
assert_eq!(failed.code.as_str(), "embed_inference_failed");
assert_eq!(serde_json::to_value(failed)?, value);
Ok(())
}
#[test]
fn initial_reader_requires_start_and_applies_fixed_limit_before_payload_read() -> TestResult {
let value = examples().remove(0).1;
let message = decode_message(Tag::Start as u8, &serde_json::to_vec(&value)?)?;
let mut bytes = Vec::new();
write_message(&mut bytes, &message, NonZeroU32::MAX)?;
assert_eq!(
read_start(&mut Cursor::new(bytes))?
.ok_or("missing Start")?
.protocol,
1
);
assert!(read_start(&mut Cursor::new([]))?.is_none());
let mut oversized = Cursor::new(8193_u32.to_be_bytes());
assert_eq!(
read_start(&mut oversized)
.err()
.ok_or("large first frame accepted")?
.code,
FailedCode::EmbedProtocol
);
assert_eq!(oversized.position(), 4);
let value = examples().remove(2).1;
let message = decode_message(Tag::Embed as u8, &serde_json::to_vec(&value)?)?;
let mut bytes = Vec::new();
write_message(&mut bytes, &message, NonZeroU32::MAX)?;
assert_eq!(
read_start(&mut Cursor::new(bytes))
.err()
.ok_or("wrong first frame accepted")?
.code,
FailedCode::EmbedProtocol
);
Ok(())
}
#[test]
fn every_payload_refuses_unknown_members_and_non_object_shapes() -> TestResult {
for (tag, mut value) in examples() {
value["unlisted"] = json!(true);
protocol(tag, &value)?;
protocol(tag, &json!([]))?;
}
Ok(())
}
#[test]
fn nested_objects_and_items_refuse_unknown_members_and_array_shapes() -> TestResult {
for (tag, mut value) in examples() {
match tag {
Tag::Start => value["limits"]["extra"] = json!(1),
Tag::Ready => value["model"]["extra"] = json!(1),
Tag::Embed | Tag::Vectors => value["items"][0]["extra"] = json!(1),
Tag::Failed => continue,
}
protocol(tag, &value)?;
}
for (tag, mut value) in examples() {
match tag {
Tag::Start => value["limits"] = json!([1024, 2, 1, 32, 16, 8192]),
Tag::Ready => value["model"] = json!([]),
Tag::Embed | Tag::Vectors => value["items"][0] = json!([]),
Tag::Failed => continue,
}
protocol(tag, &value)?;
}
Ok(())
}
#[test]
fn vector_outcomes_are_flat_and_refuse_cross_variant_members() -> TestResult {
let vector = json!({"items":[{"id":"one","outcome":"vector","vector":"0000803f","tokens_read":3,"tokens_total":3}]});
decode_message(Tag::Vectors as u8, &serde_json::to_vec(&vector)?)?;
let refused = json!({"items":[{"id":"one","outcome":"refused","code":"embed_input_empty"}]});
decode_message(Tag::Vectors as u8, &serde_json::to_vec(&refused)?)?;
let mut wrong = vector;
wrong["items"][0]["code"] = json!("embed_input_empty");
protocol(Tag::Vectors, &wrong)?;
wrong = refused;
wrong["items"][0]["vector"] = json!("0000803f");
protocol(Tag::Vectors, &wrong)?;
protocol(
Tag::Vectors,
&json!({"items":[{"id":"one","outcome":{"kind":"refused","code":"embed_input_empty"}}]}),
)?;
Ok(())
}
#[test]
fn a_refusal_without_token_counts_still_reads_and_one_with_them_reads_both() -> TestResult {
let older = json!({"items":[{"id":"one","outcome":"refused","code":"embed_input_empty"}]});
let Message::Vectors(read) = decode_message(Tag::Vectors as u8, &serde_json::to_vec(&older)?)?
else {
return Err("not vectors".into());
};
let Some(Outcome::Refused {
tokens_total: None,
tokens_limit: None,
..
}) = read.items.first()
else {
return Err("older refusal read with counts".into());
};
let longer = json!({"items":[{"id":"one","outcome":"refused","code":"text_longer_than_model","tokens_total":40,"tokens_limit":32}]});
let Message::Vectors(read) = decode_message(Tag::Vectors as u8, &serde_json::to_vec(&longer)?)?
else {
return Err("not vectors".into());
};
let Some(Outcome::Refused {
code: ItemCode::TextLongerThanModel,
tokens_total: Some(40),
tokens_limit: Some(32),
..
}) = read.items.first()
else {
return Err("counts not read".into());
};
assert_eq!(
serde_json::to_value(read.items.first())?,
longer["items"][0]
);
let older_written = serde_json::to_value(Outcome::Refused {
id: "one".to_owned(),
code: ItemCode::EmbedInputEmpty,
tokens_total: None,
tokens_limit: None,
})?;
assert_eq!(older_written, older["items"][0]);
Ok(())
}
#[test]
fn every_failure_and_item_code_has_its_exact_document_name() -> TestResult {
for (name, expected) in [
("embed_model_mismatch", FailedCode::EmbedModelMismatch),
("embed_model_missing", FailedCode::EmbedModelMissing),
("embed_limits", FailedCode::EmbedLimits),
("embed_batch_too_large", FailedCode::EmbedBatchTooLarge),
("embed_output_invalid", FailedCode::EmbedOutputInvalid),
("embed_protocol", FailedCode::EmbedProtocol),
("embed_inference_failed", FailedCode::EmbedInferenceFailed),
] {
assert_eq!(serde_json::to_value(expected)?, json!(name));
assert_eq!(serde_json::from_value::<FailedCode>(json!(name))?, expected);
}
for (name, expected) in [
("embed_input_empty", ItemCode::EmbedInputEmpty),
("embed_input_too_large", ItemCode::EmbedInputTooLarge),
("text_longer_than_model", ItemCode::TextLongerThanModel),
("windows_over_cap", ItemCode::WindowsOverCap),
] {
assert_eq!(serde_json::to_value(expected)?, json!(name));
assert_eq!(serde_json::from_value::<ItemCode>(json!(name))?, expected);
}
protocol(Tag::Failed, &json!({"code":"other","message":"bad"}))?;
protocol(
Tag::Vectors,
&json!({"items":[{"id":"one","outcome":"refused","code":"other"}]}),
)?;
Ok(())
}
#[test]
fn failed_message_limit_counts_unicode_characters_without_truncation() -> TestResult {
let valid = json!({"code":"embed_protocol","message":"é".repeat(1024)});
decode_message(Tag::Failed as u8, &serde_json::to_vec(&valid)?)?;
protocol(
Tag::Failed,
&json!({"code":"embed_protocol","message":"é".repeat(1025)}),
)?;
let too_long = Message::Failed(Failed {
code: FailedCode::EmbedProtocol,
message: "x".repeat(1025),
});
let mut bytes = Vec::new();
let error = write_message(&mut bytes, &too_long, NonZeroU32::MAX)
.err()
.ok_or("long failure emitted")?;
assert_eq!(error.code, FailedCode::EmbedProtocol);
assert!(bytes.is_empty());
Ok(())
}
#[test]
fn numeric_members_are_whole_and_vector_members_are_text() -> TestResult {
for (tag, mut value) in examples() {
match tag {
Tag::Start => value["limits"]["tokens"] = json!(1.5),
Tag::Ready => value["model"]["dimensions"] = json!(1.5),
Tag::Vectors => value["items"][0]["tokens_read"] = json!(1.5),
_ => continue,
}
protocol(tag, &value)?;
}
protocol(
Tag::Vectors,
&json!({"items":[{"id":"one","outcome":"vector","vector":[1.0],"tokens_read":1,"tokens_total":1}]}),
)?;
let value = json!({"items":[{"id":"one","outcome":"vector","vector":"0000803f","tokens_read":u64::MAX,"tokens_total":u64::MAX}]});
decode_message(Tag::Vectors as u8, &serde_json::to_vec(&value)?)?;
Ok(())
}
#[test]
fn unknown_distance_and_kind_are_protocol_refusals() -> TestResult {
let mut value = examples().remove(1).1;
value["model"]["distance"] = json!("manhattan");
protocol(Tag::Ready, &value)?;
protocol(
Tag::Embed,
&json!({"kind":"other","items":[{"id":"one","text":"text"}]}),
)?;
Ok(())
}
#[test]
fn clean_eof_and_bad_frame_have_distinct_results() -> TestResult {
assert!(read_message(&mut Cursor::new([]), NonZeroU32::MAX)?.is_none());
for bytes in [b"\0".as_slice(), b"\0\0\0\x02x", b"\0\0\0\0"] {
let error = read_message(&mut Cursor::new(bytes), NonZeroU32::MAX)
.err()
.ok_or("bad frame accepted")?;
assert_eq!(error.code, FailedCode::EmbedProtocol);
}
Ok(())
}
#[test]
fn start_fixed_limit_and_later_frame_limit_are_protocol_refusals() -> TestResult {
let mut value = examples().remove(0).1;
value["model_dir"] = json!(format!("/{}", "m".repeat(8192)));
protocol(Tag::Start, &value)?;
let message = Message::Embed(Embed {
kind: Kind::Query,
items: vec![Input {
id: "one".to_owned(),
text: "text".to_owned(),
}],
});
let mut output = Vec::new();
let error = write_message(&mut output, &message, NonZeroU32::MIN)
.err()
.ok_or("overlimit frame emitted")?;
assert_eq!(error.code, FailedCode::EmbedProtocol);
assert!(output.is_empty());
let mut input = Cursor::new(b"\0\0\0\x02ab");
let error = read_message(&mut input, NonZeroU32::MIN)
.err()
.ok_or("overlimit frame read")?;
assert_eq!(error.code, FailedCode::EmbedProtocol);
assert_eq!(input.position(), 4);
Ok(())
}