use rtx_llm_tools::templates::{PromptTemplate, TemplateContext, TemplateEngine}; use serde_json::Value; #[tokio::test] async fn test_simple_variable_substitution() { let template_str = "Hello {name}, welcome to {system}!"; let template = PromptTemplate::from_str(template_str).unwrap(); let mut context = TemplateContext::new(); context.insert("name", Value::String("Alice".to_string())); context.insert("system", Value::String("RustyTorch".to_string())); let result = template.render(&context).await; assert!(result.is_ok()); assert_eq!(result.unwrap(), "Hello Alice, welcome to RustyTorch!"); } #[tokio::test] async fn test_conditional_rendering() { let template_str = r#" {%- if is_admin -%} Welcome admin {name}! You have full access. {%- else -%} Hello {name}, you have limited access. {%- endif -%} "#; let template = PromptTemplate::from_str(template_str).unwrap(); let mut admin_context = TemplateContext::new(); admin_context.insert("name", Value::String("Bob".to_string())); admin_context.insert("is_admin", Value::Bool(true)); let result = template.render(&admin_context).await.unwrap(); assert!(result.contains("Welcome admin Bob! You have full access.")); let mut user_context = TemplateContext::new(); user_context.insert("name", Value::String("Charlie".to_string())); user_context.insert("is_admin", Value::Bool(false)); let result = template.render(&user_context).await.unwrap(); assert!(result.contains("Hello Charlie, you have limited access.")); } #[tokio::test] async fn test_loop_iteration() { let template_str = r#" Items: {%- for item in items %} - {item.name}: {item.price} {%- endfor -%} "#; let template = PromptTemplate::from_str(template_str).unwrap(); let items = vec![ serde_json::json!({"name": "Apple", "price": "$1.50"}), serde_json::json!({"name": "Banana", "price": "$0.75"}), ]; let mut context = TemplateContext::new(); context.insert("items", Value::Array(items)); let result = template.render(&context).await.unwrap(); assert!(result.contains("- Apple: $1.50")); assert!(result.contains("- Banana: $0.75")); } #[tokio::test] async fn test_nested_object_access() { let template_str = "User: {user.profile.name}, Age: {user.profile.age}"; let template = PromptTemplate::from_str(template_str).unwrap(); let user_data = serde_json::json!({ "user": { "profile": { "name": "David", "age": 30 } } }); let mut context = TemplateContext::new(); context.insert_nested(user_data); let result = template.render(&context).await.unwrap(); assert_eq!(result, "User: David, Age: 30"); } #[tokio::test] async fn test_template_inheritance() { let base_template = r#" System: {system_name} {%- block content -%} Default content {%- endblock -%} "#; let child_template = r#" {%- extends base -%} {%- block content -%} Task: {task_description} Parameters: {params} {%- endblock -%} "#; let mut engine = TemplateEngine::new(); engine .register_template("base", base_template) .await .unwrap(); let template = engine.compile_template(child_template).await.unwrap(); let mut context = TemplateContext::new(); context.insert("system_name", Value::String("RustyTorch".to_string())); context.insert( "task_description", Value::String("Generate text".to_string()), ); context.insert("params", Value::String("{\"max_tokens\": 100}".to_string())); let result = template.render(&context).await.unwrap(); assert!(result.contains("System: RustyTorch")); assert!(result.contains("Task: Generate text")); assert!(result.contains("Parameters: {\"max_tokens\": 100}")); } #[tokio::test] async fn test_custom_functions() { let template_str = "Current time: {now() | format_time('%Y-%m-%d %H:%M')}"; let template = PromptTemplate::from_str(template_str).unwrap(); let mut context = TemplateContext::new(); context.register_function( "now", Box::new(|| Ok(Value::String(chrono::Utc::now().to_rfc3339()))), ); context.register_filter( "format_time", Box::new(|_value: &Value, _format: &str| { // Mock implementation for test Ok(Value::String("2024-01-15 10:30".to_string())) }), ); let result = template.render(&context).await.unwrap(); assert!(result.contains("Current time: 2024-01-15 10:30")); } #[tokio::test] async fn test_error_handling_missing_variable() { let template_str = "Hello {missing_name}!"; let template = PromptTemplate::from_str(template_str).unwrap(); let context = TemplateContext::new(); let result = template.render(&context).await; assert!(result.is_err()); let error = result.unwrap_err(); assert!( error.to_string().contains("missing_name"), "Expected MissingVariable error but got: {}", error ); } #[tokio::test] async fn test_template_security_injection_prevention() { let malicious_input = "'; DROP TABLE users; --"; let template_str = "Query: {user_input}"; let template = PromptTemplate::from_str(template_str).unwrap(); let mut context = TemplateContext::new(); context.insert("user_input", Value::String(malicious_input.to_string())); let result = template.render(&context).await.unwrap(); // Should escape dangerous characters assert!(result.contains("Query: '; DROP TABLE users; --")); } #[tokio::test] async fn test_template_caching() { let template_str = "Hello {name}!"; let mut engine = TemplateEngine::new(); engine.enable_caching(true); // First compilation should cache the template let template1 = engine.compile_template(template_str).await.unwrap(); let template2 = engine.compile_template(template_str).await.unwrap(); // Should use cached version (test by checking compilation time or cache stats) assert_eq!(template1.template_id(), template2.template_id()); } #[tokio::test] async fn test_template_validation() { let invalid_template = "Hello {unclosed_brace"; let result = PromptTemplate::from_str(invalid_template); assert!(result.is_err()); let error = result.unwrap_err(); assert!( error.to_string().contains("unclosed") || error.to_string().contains("syntax"), "Expected SyntaxError but got: {}", error ); } #[tokio::test] async fn test_partial_template_application() { let template_str = "System: {system}, Task: {task}, User: {user}"; let template = PromptTemplate::from_str(template_str).unwrap(); let mut partial_context = TemplateContext::new(); partial_context.insert("system", Value::String("RustyTorch".to_string())); partial_context.insert("task", Value::String("Generate".to_string())); let partial_template = template.bind_partial(&partial_context).await.unwrap(); let mut final_context = TemplateContext::new(); final_context.insert("user", Value::String("Alice".to_string())); let result = partial_template.render(&final_context).await.unwrap(); assert_eq!(result, "System: RustyTorch, Task: Generate, User: Alice"); }