use crate::fluent::Tool;
use anyhow::{Context, Result};
use serde::Deserialize;
use std::collections::HashMap;
use std::fs;
use std::path::Path;
#[derive(Debug, thiserror::Error)]
pub enum YamlError {
#[error("Unknown tool: '{0}'. Available tools: Read, Write, Edit, Bash, Glob, Grep, Task, WebFetch, WebSearch, NotebookEdit, AskUserQuestion, TodoWrite, Skill")]
UnknownTool(String),
#[error("IO error: {0}")]
Io(#[from] std::io::Error),
#[error("YAML parse error: {0}")]
Yaml(#[from] serde_yaml::Error),
}
#[derive(Debug, Deserialize)]
pub struct Test {
pub name: String,
pub prompt: String,
#[serde(default)]
pub agent: Option<String>,
pub assertions: Vec<Assertion>,
}
#[derive(Debug, Deserialize)]
pub struct Assertion {
pub tool: Option<String>,
#[serde(default = "default_true")]
pub called: bool,
pub params: Option<HashMap<String, String>>,
pub called_after: Option<String>,
pub called_before: Option<String>,
pub call_count: Option<u32>,
pub max_calls: Option<u32>,
pub min_calls: Option<u32>,
pub nth_call_params: Option<HashMap<u32, HashMap<String, String>>>,
pub first_call_params: Option<HashMap<String, String>>,
pub last_call_params: Option<HashMap<String, String>>,
pub stdout: Option<StdoutConstraints>,
}
#[derive(Debug, Deserialize, Clone)]
pub struct StdoutConstraints {
#[serde(default = "default_true")]
pub exists: bool,
pub contains: Option<String>,
pub not_contains: Option<String>,
pub matches: Option<String>,
pub not_matches: Option<String>,
}
fn default_true() -> bool {
true
}
pub fn load_test(path: &Path) -> Result<Test> {
let content = fs::read_to_string(path).context("Failed to read test file")?;
let test: Test = serde_yaml::from_str(&content).context("Failed to parse YAML")?;
Ok(test)
}
pub fn parse_tool_name(s: &str) -> Result<Tool, YamlError> {
match s.to_lowercase().as_str() {
"read" => Ok(Tool::Read),
"write" => Ok(Tool::Write),
"edit" => Ok(Tool::Edit),
"bash" => Ok(Tool::Bash),
"glob" => Ok(Tool::Glob),
"grep" => Ok(Tool::Grep),
"task" => Ok(Tool::Task),
"webfetch" => Ok(Tool::WebFetch),
"websearch" => Ok(Tool::WebSearch),
"notebookedit" => Ok(Tool::NotebookEdit),
"askuserquestion" => Ok(Tool::AskUserQuestion),
"todowrite" => Ok(Tool::TodoWrite),
"killshell" => Ok(Tool::KillShell),
"taskoutput" => Ok(Tool::TaskOutput),
"skill" => Ok(Tool::Skill),
"read_file" => Ok(Tool::Read),
"write_file" => Ok(Tool::Write),
"edit_file" => Ok(Tool::Edit),
"execute_command" => Ok(Tool::Bash),
"glob_files" => Ok(Tool::Glob),
"search_files" => Ok(Tool::Grep),
"web_fetch" => Ok(Tool::WebFetch),
"web_search" => Ok(Tool::WebSearch),
"notebook_edit" => Ok(Tool::NotebookEdit),
"ask_user" | "ask_user_question" => Ok(Tool::AskUserQuestion),
"todo_write" => Ok(Tool::TodoWrite),
"kill_shell" => Ok(Tool::KillShell),
"task_output" => Ok(Tool::TaskOutput),
_ => Err(YamlError::UnknownTool(s.to_string())),
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse_tool_name_primary() {
assert_eq!(parse_tool_name("Read").unwrap(), Tool::Read);
assert_eq!(parse_tool_name("Write").unwrap(), Tool::Write);
assert_eq!(parse_tool_name("Bash").unwrap(), Tool::Bash);
assert_eq!(parse_tool_name("WebFetch").unwrap(), Tool::WebFetch);
}
#[test]
fn test_parse_tool_name_case_insensitive() {
assert_eq!(parse_tool_name("read").unwrap(), Tool::Read);
assert_eq!(parse_tool_name("READ").unwrap(), Tool::Read);
assert_eq!(parse_tool_name("ReAd").unwrap(), Tool::Read);
}
#[test]
fn test_parse_tool_name_aliases() {
assert_eq!(parse_tool_name("read_file").unwrap(), Tool::Read);
assert_eq!(parse_tool_name("write_file").unwrap(), Tool::Write);
assert_eq!(parse_tool_name("execute_command").unwrap(), Tool::Bash);
assert_eq!(parse_tool_name("search_files").unwrap(), Tool::Grep);
}
#[test]
fn test_parse_tool_name_unknown() {
assert!(parse_tool_name("unknown_tool").is_err());
assert!(parse_tool_name("").is_err());
}
#[test]
fn test_deserialize_assertion() {
let yaml = r#"
tool: Read
called: true
params:
file_path: "*.txt"
"#;
let assertion: Assertion = serde_yaml::from_str(yaml).unwrap();
assert_eq!(assertion.tool, Some("Read".to_string()));
assert!(assertion.called);
assert!(assertion.params.is_some());
}
#[test]
fn test_deserialize_stdout_assertion() {
let yaml = r#"
stdout:
exists: true
contains: "success"
not_contains: "error"
"#;
let assertion: Assertion = serde_yaml::from_str(yaml).unwrap();
assert!(assertion.tool.is_none());
assert!(assertion.stdout.is_some());
let stdout = assertion.stdout.unwrap();
assert!(stdout.exists);
assert_eq!(stdout.contains, Some("success".to_string()));
assert_eq!(stdout.not_contains, Some("error".to_string()));
}
#[test]
fn test_deserialize_test() {
let yaml = r#"
name: "Test reading files"
prompt: "Read the config"
assertions:
- tool: Read
called: true
"#;
let test: Test = serde_yaml::from_str(yaml).unwrap();
assert_eq!(test.name, "Test reading files");
assert_eq!(test.prompt, "Read the config");
assert_eq!(test.assertions.len(), 1);
}
#[test]
fn test_default_called_true() {
let yaml = r#"
tool: Read
"#;
let assertion: Assertion = serde_yaml::from_str(yaml).unwrap();
assert!(assertion.called);
}
}