use simplify_baml::*;
use simplify_baml_macros::BamlSchema;
use std::collections::HashMap;
#[derive(BamlSchema)]
#[allow(dead_code)]
struct Config {
settings: HashMap<String, String>,
}
#[derive(BamlSchema)]
#[allow(dead_code)]
struct UserScores {
scores: HashMap<String, i64>,
}
#[test]
fn test_hashmap_string_string_schema() {
let registry = BamlSchemaRegistry::new()
.register::<Config>();
let ir = registry.build();
let config_class = ir.classes.iter().find(|c| c.name == "Config").unwrap();
let settings_field = config_class.fields.iter().find(|f| f.name == "settings").unwrap();
assert_eq!(
settings_field.field_type,
FieldType::Map(Box::new(FieldType::String), Box::new(FieldType::String))
);
}
#[test]
fn test_hashmap_string_int_schema() {
let registry = BamlSchemaRegistry::new()
.register::<UserScores>();
let ir = registry.build();
let class = ir.classes.iter().find(|c| c.name == "UserScores").unwrap();
let field = class.fields.iter().find(|f| f.name == "scores").unwrap();
assert_eq!(
field.field_type,
FieldType::Map(Box::new(FieldType::String), Box::new(FieldType::Int))
);
}
#[test]
fn test_parse_hashmap_response() {
let mut ir = IR::new();
ir.classes.push(Class {
name: "Config".to_string(),
description: None,
fields: vec![Field {
name: "settings".to_string(),
field_type: FieldType::Map(Box::new(FieldType::String), Box::new(FieldType::String)),
optional: false,
description: None,
}],
});
let response = r#"{"settings": {"theme": "dark", "language": "en"}}"#;
let result = parse_llm_response_with_ir(
&ir,
response,
&FieldType::Class("Config".to_string()),
);
assert!(result.is_ok());
if let Ok(BamlValue::Map(map)) = result {
if let Some(BamlValue::Map(settings)) = map.get("settings") {
match settings.get("theme") {
Some(BamlValue::String(v)) => assert_eq!(v, "dark"),
other => panic!("Expected String 'dark', got {:?}", other),
}
match settings.get("language") {
Some(BamlValue::String(v)) => assert_eq!(v, "en"),
other => panic!("Expected String 'en', got {:?}", other),
}
} else {
panic!("Expected settings to be a map");
}
}
}