test: Added unit tests for the rag, completions and prompt, macros, vault, and functions/tool usage

This commit is contained in:
2026-05-01 13:24:58 -06:00
parent e23e5f9f7b
commit 349b3748bd
12 changed files with 801 additions and 46 deletions
+202
View File
@@ -168,3 +168,205 @@ pub struct MacroVariable {
pub rest: bool,
pub default: Option<String>,
}
#[cfg(test)]
mod tests {
use super::*;
fn var(name: &str, rest: bool, default: Option<&str>) -> MacroVariable {
MacroVariable {
name: name.to_string(),
rest,
default: default.map(String::from),
}
}
fn macro_with_vars(vars: Vec<MacroVariable>) -> Macro {
Macro {
variables: vars,
steps: vec![],
}
}
#[test]
fn resolve_no_variables() {
let m = macro_with_vars(vec![]);
let result = m.resolve_variables(&[]).unwrap();
assert!(result.is_empty());
}
#[test]
fn resolve_required_variable_provided() {
let m = macro_with_vars(vec![var("name", false, None)]);
let result = m.resolve_variables(&["Alice".into()]).unwrap();
assert_eq!(result["name"], "Alice");
}
#[test]
fn resolve_required_variable_missing_errors() {
let m = macro_with_vars(vec![var("name", false, None)]);
let result = m.resolve_variables(&[]);
assert!(result.is_err());
assert!(result.unwrap_err().to_string().contains("name"));
}
#[test]
fn resolve_default_variable_uses_default() {
let m = macro_with_vars(vec![var("color", false, Some("blue"))]);
let result = m.resolve_variables(&[]).unwrap();
assert_eq!(result["color"], "blue");
}
#[test]
fn resolve_default_variable_overridden() {
let m = macro_with_vars(vec![var("color", false, Some("blue"))]);
let result = m.resolve_variables(&["red".into()]).unwrap();
assert_eq!(result["color"], "red");
}
#[test]
fn resolve_rest_variable_captures_all_remaining() {
let m = macro_with_vars(vec![var("first", false, None), var("rest", true, None)]);
let result = m
.resolve_variables(&["a".into(), "b".into(), "c".into()])
.unwrap();
assert_eq!(result["first"], "a");
assert_eq!(result["rest"], "b c");
}
#[test]
fn resolve_rest_variable_with_default() {
let m = macro_with_vars(vec![var("args", true, Some("default text"))]);
let result = m.resolve_variables(&[]).unwrap();
assert_eq!(result["args"], "default text");
}
#[test]
fn resolve_multiple_variables() {
let m = macro_with_vars(vec![
var("a", false, None),
var("b", false, None),
var("c", false, Some("default_c")),
]);
let result = m.resolve_variables(&["x".into(), "y".into()]).unwrap();
assert_eq!(result["a"], "x");
assert_eq!(result["b"], "y");
assert_eq!(result["c"], "default_c");
}
#[test]
fn usage_no_variables() {
let m = macro_with_vars(vec![]);
assert_eq!(m.usage("my-macro"), "my-macro");
}
#[test]
fn usage_required_variable() {
let m = macro_with_vars(vec![var("name", false, None)]);
assert_eq!(m.usage("greet"), "greet <name>");
}
#[test]
fn usage_optional_variable() {
let m = macro_with_vars(vec![var("color", false, Some("blue"))]);
assert_eq!(m.usage("paint"), "paint [color]");
}
#[test]
fn usage_rest_variable() {
let m = macro_with_vars(vec![var("args", true, None)]);
assert_eq!(m.usage("run"), "run <args>...");
}
#[test]
fn usage_rest_with_default() {
let m = macro_with_vars(vec![var("args", true, Some("default"))]);
assert_eq!(m.usage("run"), "run [args]...");
}
#[test]
fn usage_mixed_variables() {
let m = macro_with_vars(vec![
var("target", false, None),
var("flags", true, Some("")),
]);
assert_eq!(m.usage("build"), "build <target> [flags]...");
}
#[test]
fn interpolate_replaces_variables() {
let vars = IndexMap::from([("name".to_string(), "world".to_string())]);
let result = Macro::interpolate_command("hello {{name}}", &vars);
assert_eq!(result, "hello world");
}
#[test]
fn interpolate_multiple_variables() {
let vars = IndexMap::from([
("a".to_string(), "1".to_string()),
("b".to_string(), "2".to_string()),
]);
let result = Macro::interpolate_command("{{a}} + {{b}}", &vars);
assert_eq!(result, "1 + 2");
}
#[test]
fn interpolate_no_variables_passthrough() {
let vars = IndexMap::new();
let result = Macro::interpolate_command("no vars here", &vars);
assert_eq!(result, "no vars here");
}
#[test]
fn interpolate_variable_not_found_left_as_is() {
let vars = IndexMap::new();
let result = Macro::interpolate_command("hello {{missing}}", &vars);
assert_eq!(result, "hello {{missing}}");
}
#[test]
fn deserialize_macro_from_yaml() {
let yaml = r#"
steps:
- ".role coder"
- "write code for {{task}}"
variables:
- name: task
"#;
let m: Macro = serde_yaml::from_str(yaml).unwrap();
assert_eq!(m.steps.len(), 2);
assert_eq!(m.variables.len(), 1);
assert_eq!(m.variables[0].name, "task");
assert!(!m.variables[0].rest);
assert!(m.variables[0].default.is_none());
}
#[test]
fn deserialize_macro_with_defaults() {
let yaml = r#"
steps:
- "test"
variables:
- name: mode
default: "fast"
- name: args
rest: true
default: "none"
"#;
let m: Macro = serde_yaml::from_str(yaml).unwrap();
assert_eq!(m.variables[0].default, Some("fast".to_string()));
assert!(m.variables[1].rest);
assert_eq!(m.variables[1].default, Some("none".to_string()));
}
#[test]
fn deserialize_macro_no_variables() {
let yaml = r#"
steps:
- ".help"
"#;
let m: Macro = serde_yaml::from_str(yaml).unwrap();
assert!(m.variables.is_empty());
assert_eq!(m.steps.len(), 1);
}
}
+45
View File
@@ -234,3 +234,48 @@ pub(crate) fn generate_declarations<L: ScriptedLanguage>(
}
Ok(out)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn underscore_simple() {
assert_eq!(underscore("hello"), "hello");
}
#[test]
fn underscore_dashes_to_underscores() {
assert_eq!(underscore("my-func-name"), "my_func_name");
}
#[test]
fn underscore_spaces_to_underscores() {
assert_eq!(underscore("my func"), "my_func");
}
#[test]
fn underscore_special_chars_removed() {
assert_eq!(underscore("func@name!here"), "func_name_here");
}
#[test]
fn underscore_consecutive_specials_collapsed() {
assert_eq!(underscore("a---b"), "a_b");
}
#[test]
fn underscore_leading_trailing_stripped() {
assert_eq!(underscore("-leading-"), "leading");
}
#[test]
fn underscore_uppercase_lowered() {
assert_eq!(underscore("MyFunc"), "myfunc");
}
#[test]
fn underscore_mixed() {
assert_eq!(underscore("Get-User Info"), "get_user_info");
}
}
+234
View File
@@ -1080,3 +1080,237 @@ fn reciprocal_rank_fusion(
.map(|(v, _)| v)
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn document_id_round_trip() {
let id = DocumentId::new(5, 17);
let (file, doc) = id.split();
assert_eq!(file, 5);
assert_eq!(doc, 17);
}
#[test]
fn document_id_zero_zero() {
let id = DocumentId::new(0, 0);
let (file, doc) = id.split();
assert_eq!(file, 0);
assert_eq!(doc, 0);
}
#[test]
fn document_id_large_values() {
let id = DocumentId::new(1000, 9999);
let (file, doc) = id.split();
assert_eq!(file, 1000);
assert_eq!(doc, 9999);
}
#[test]
fn document_id_debug_format() {
let id = DocumentId::new(3, 7);
let formatted = format!("{id:?}");
assert_eq!(formatted, "3-7");
}
#[test]
fn document_id_equality() {
let a = DocumentId::new(1, 2);
let b = DocumentId::new(1, 2);
assert_eq!(a, b);
}
#[test]
fn document_id_inequality() {
let a = DocumentId::new(1, 2);
let b = DocumentId::new(1, 3);
assert_ne!(a, b);
}
#[test]
fn document_id_ordering() {
let a = DocumentId::new(0, 1);
let b = DocumentId::new(1, 0);
assert!(a < b);
}
#[test]
fn rag_document_new() {
let doc = RagDocument::new("hello world");
assert_eq!(doc.page_content, "hello world");
assert!(doc.metadata.is_empty());
}
#[test]
fn rag_document_default() {
let doc = RagDocument::default();
assert_eq!(doc.page_content, "");
assert!(doc.metadata.is_empty());
}
#[test]
fn rag_data_new_defaults() {
let data = RagData::new("model".into(), 1000, 20, None, 5, None);
assert_eq!(data.embedding_model, "model");
assert_eq!(data.chunk_size, 1000);
assert_eq!(data.chunk_overlap, 20);
assert_eq!(data.top_k, 5);
assert!(data.reranker_model.is_none());
assert!(data.files.is_empty());
assert!(data.vectors.is_empty());
assert!(data.document_paths.is_empty());
assert_eq!(data.next_file_id, 0);
}
#[test]
fn rag_data_get_returns_document() {
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
let file = RagFile {
hash: "abc".into(),
path: "test.txt".into(),
documents: vec![RagDocument::new("first"), RagDocument::new("second")],
};
data.files.insert(0, file);
let doc = data.get(DocumentId::new(0, 0)).unwrap();
assert_eq!(doc.page_content, "first");
let doc = data.get(DocumentId::new(0, 1)).unwrap();
assert_eq!(doc.page_content, "second");
}
#[test]
fn rag_data_get_returns_none_for_missing_file() {
let data = RagData::new("m".into(), 100, 10, None, 5, None);
assert!(data.get(DocumentId::new(99, 0)).is_none());
}
#[test]
fn rag_data_get_returns_none_for_missing_document() {
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
let file = RagFile {
hash: "abc".into(),
path: "test.txt".into(),
documents: vec![RagDocument::new("only one")],
};
data.files.insert(0, file);
assert!(data.get(DocumentId::new(0, 5)).is_none());
}
#[test]
fn rag_data_del_removes_files_and_vectors() {
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
let file = RagFile {
hash: "abc".into(),
path: "test.txt".into(),
documents: vec![RagDocument::new("doc")],
};
data.files.insert(0, file);
let doc_id = DocumentId::new(0, 0);
data.vectors.insert(doc_id, vec![0.1, 0.2, 0.3]);
assert!(data.files.contains_key(&0));
assert!(data.vectors.contains_key(&doc_id));
data.del(vec![0]);
assert!(!data.files.contains_key(&0));
assert!(!data.vectors.contains_key(&doc_id));
}
#[test]
fn rag_data_del_nonexistent_is_noop() {
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
data.del(vec![99]);
assert!(data.files.is_empty());
}
#[test]
fn rag_data_add_inserts_files_and_vectors() {
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
let file = RagFile {
hash: "xyz".into(),
path: "new.txt".into(),
documents: vec![RagDocument::new("content")],
};
let doc_id = DocumentId::new(0, 0);
let embeddings = vec![vec![0.5, 0.6, 0.7]];
data.add(1, vec![(0, file)], vec![doc_id], embeddings);
assert_eq!(data.next_file_id, 1);
assert!(data.files.contains_key(&0));
assert!(data.vectors.contains_key(&doc_id));
assert_eq!(data.vectors[&doc_id], vec![0.5, 0.6, 0.7]);
}
#[test]
fn rag_template_contains_placeholders() {
assert!(RAG_TEMPLATE.contains("__CONTEXT__"));
assert!(RAG_TEMPLATE.contains("__SOURCES__"));
assert!(RAG_TEMPLATE.contains("__INPUT__"));
}
#[test]
fn get_separators_returns_language_specific() {
let rs_seps = splitter::get_separators("rs");
assert!(rs_seps.iter().any(|s| s.contains("fn ")));
let py_seps = splitter::get_separators("py");
assert!(py_seps.iter().any(|s| s.contains("def ")));
let md_seps = splitter::get_separators("md");
assert!(md_seps.iter().any(|s| s.contains("# ")));
}
#[test]
fn get_separators_unknown_returns_defaults() {
let seps = get_separators("xyz");
assert_eq!(seps, DEFAULT_SEPARATORS.to_vec());
}
#[test]
fn get_separators_all_known_extensions() {
let known = [
"c", "cc", "cpp", "go", "java", "js", "mjs", "cjs", "php", "proto", "py", "rst", "rb",
"rs", "scala", "swift", "md", "mkd", "tex", "htm", "html", "sol",
];
for ext in known {
let seps = get_separators(ext);
assert_ne!(seps, DEFAULT_SEPARATORS.to_vec(), "Extension '{ext}' should have language-specific separators");
}
}
#[test]
fn rag_data_build_bm25_empty() {
let data = RagData::new("m".into(), 100, 10, None, 5, None);
let engine = data.build_bm25();
let results = engine.search("anything", 5);
assert!(results.is_empty());
}
#[test]
fn rag_data_build_bm25_finds_documents() {
let mut data = RagData::new("m".into(), 100, 10, None, 5, None);
let file = RagFile {
hash: "h".into(),
path: "test.txt".into(),
documents: vec![
RagDocument::new("rust programming language"),
RagDocument::new("python scripting language"),
],
};
data.files.insert(0, file);
let engine = data.build_bm25();
let results = engine.search("rust", 5);
assert!(!results.is_empty());
let top = &results[0];
let (file_idx, doc_idx) = top.document.id.split();
assert_eq!(file_idx, 0);
assert_eq!(doc_idx, 0);
}
}
+42
View File
@@ -154,3 +154,45 @@ impl Vault {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn secret_re_matches_double_braces() {
let captures = SECRET_RE.captures("{{MY_SECRET}}").unwrap().unwrap();
assert_eq!(&captures[1], "MY_SECRET");
}
#[test]
fn secret_re_matches_with_surrounding_text() {
let text = "key={{API_KEY}} here";
let captures = SECRET_RE.captures(text).unwrap().unwrap();
assert_eq!(&captures[1], "API_KEY");
}
#[test]
fn secret_re_no_match_single_braces() {
let result = SECRET_RE.captures("{NOT_SECRET}").unwrap();
assert!(result.is_none());
}
#[test]
fn secret_re_no_match_plain_text() {
let result = SECRET_RE.captures("just plain text").unwrap();
assert!(result.is_none());
}
#[test]
fn secret_re_matches_with_spaces() {
let captures = SECRET_RE.captures("{{ SPACED }}").unwrap().unwrap();
assert_eq!(&captures[1], " SPACED ");
}
#[test]
fn vault_default_creates_instance() {
let vault = Vault::default();
assert!(vault.password_file().is_err());
}
}