use embellama::{
BackendInfo, BackendType, EmbeddingEngine, EngineConfig, detect_best_backend,
get_compiled_backend,
};
use std::env;
fn get_test_model_path() -> String {
env::var("EMBELLAMA_TEST_MODEL").unwrap_or_else(|_| {
panic!(
"EMBELLAMA_TEST_MODEL environment variable not set.\n\
Please run: just download-test-model"
)
})
}
#[test]
fn test_backend_detection() {
let backend = detect_best_backend();
let compiled = get_compiled_backend();
println!("Detected backend: {:?}", backend);
println!("Compiled backend: {:?}", compiled);
assert_eq!(backend, compiled);
match backend {
BackendType::Cpu
| BackendType::OpenMP
| BackendType::ROCm
| BackendType::Cuda
| BackendType::Metal
| BackendType::Vulkan => {
}
}
}
#[test]
fn test_backend_info() {
let info = BackendInfo::new();
println!("Backend info:\n{}", info);
assert!(!info.platform.is_empty());
assert_eq!(info.backend, info.compiled_backend);
#[cfg(feature = "openmp")]
assert!(info.available_features.contains(&"openmp".to_string()));
#[cfg(feature = "metal")]
assert!(info.available_features.contains(&"metal".to_string()));
#[cfg(feature = "cuda")]
assert!(info.available_features.contains(&"cuda".to_string()));
#[cfg(feature = "vulkan")]
assert!(info.available_features.contains(&"vulkan".to_string()));
}
#[test]
fn test_backend_gpu_acceleration() {
let backend = detect_best_backend();
#[cfg(any(feature = "metal", feature = "cuda", feature = "vulkan"))]
{
if matches!(
backend,
BackendType::Metal | BackendType::Cuda | BackendType::Vulkan
) {
assert!(backend.is_gpu_accelerated());
assert_eq!(backend.recommended_gpu_layers(), Some(999));
}
}
if matches!(backend, BackendType::Cpu | BackendType::OpenMP) {
assert!(!backend.is_gpu_accelerated());
assert_eq!(backend.recommended_gpu_layers(), None);
}
}
#[test]
fn test_config_with_backend_detection() {
use tempfile::NamedTempFile;
let temp_file = NamedTempFile::new().expect("Failed to create temp file");
let model_path = temp_file.path().with_extension("gguf");
std::fs::write(&model_path, b"GGUF").expect("Failed to write temp file");
let config = EngineConfig::with_backend_detection()
.with_model_path(model_path.to_str().unwrap())
.with_model_name("test-backend")
.build()
.expect("Failed to build config");
let backend = detect_best_backend();
if backend.is_gpu_accelerated() {
assert!(config.use_gpu);
assert_eq!(config.model_config.n_gpu_layers, Some(999));
}
config.validate().expect("Config validation failed");
}
#[test]
#[cfg(feature = "metal")]
fn test_metal_backend_on_macos() {
#[cfg(target_os = "macos")]
{
let backend = detect_best_backend();
assert_eq!(backend, BackendType::Metal);
assert!(backend.is_gpu_accelerated());
}
}
#[test]
fn test_backend_display() {
assert_eq!(BackendType::Cpu.to_string(), "CPU");
assert_eq!(BackendType::OpenMP.to_string(), "OpenMP");
assert_eq!(BackendType::ROCm.to_string(), "ROCm");
assert_eq!(BackendType::Cuda.to_string(), "CUDA");
assert_eq!(BackendType::Metal.to_string(), "Metal");
assert_eq!(BackendType::Vulkan.to_string(), "Vulkan");
}
#[test]
#[ignore = "Requires model file"]
fn test_engine_with_backend_detection() {
let model_path = get_test_model_path();
let config = EngineConfig::with_backend_detection()
.with_model_path(&model_path)
.with_model_name("test-backend")
.with_context_size(2048) .build()
.expect("Failed to build config");
let engine = EmbeddingEngine::new(config).expect("Failed to create engine");
let text = "Test embedding with backend detection";
let embedding = engine
.embed(Some("test-backend"), text)
.expect("Failed to generate embedding");
assert!(!embedding.is_empty());
println!("Generated embedding with {} dimensions", embedding.len());
let info = BackendInfo::new();
println!("Used backend: {}", info.backend);
}