Newer
Older
TelosDB / src / backend / tests / vector_test.rs
use app_lib::db;
use sea_orm::{ConnectionTrait, DatabaseBackend, Statement};
use std::env;
use std::path::PathBuf;

#[tokio::test]
async fn test_vec0_extension_loading() {
    // 1. Determine paths
    let test_db = "test-vector-vec0.db";
    if std::path::Path::new(test_db).exists() {
        let _ = std::fs::remove_file(test_db);
    }

    let mut candidates = vec![
        PathBuf::from("bin/vec0.dll"),
        PathBuf::from("../bin/vec0.dll"),
        PathBuf::from("../../bin/vec0.dll"),
        env::current_dir().unwrap().join("bin").join("vec0.dll"),
        env::current_dir().unwrap().join("../").join("bin").join("vec0.dll"),
        env::current_dir().unwrap().join("../../").join("bin").join("vec0.dll"),
    ];

    let mut ext_path = PathBuf::from("vec0.dll");
    for cand in candidates {
        if cand.exists() {
            ext_path = cand;
            break;
        }
    }

    let ext_path_str = ext_path.to_string_lossy().to_string();
    println!("Testing with extension path: {}", ext_path_str);

    // 2. Initialize DB
    let conn = db::init_db(test_db, &ext_path_str).await.expect("Failed to initialize DB with vec0");

    // 3. Verify vec0 is actually working by querying the virtual table
    
    // Test inserting a vector
    let embedding: Vec<f32> = vec![0.1; 384];
    let embedding_bytes: Vec<u8> = embedding.iter().flat_map(|f| f.to_le_bytes()).collect();

    let insert_res = conn.execute(Statement::from_sql_and_values(
        DatabaseBackend::Sqlite,
        "INSERT INTO vec_items (id, embedding) VALUES (?, ?)",
        [1.into(), embedding_bytes.clone().into()],
    )).await;

    assert!(insert_res.is_ok(), "Failed to insert into vec_items: {:?}", insert_res.err());

    // Test a simple vector search (MATCH)
    // Using Statement directly with query_all
    let search_res = conn.query_all(Statement::from_sql_and_values(
        DatabaseBackend::Sqlite,
        "SELECT id, distance FROM vec_items WHERE embedding MATCH ? ORDER BY distance LIMIT 1",
        [embedding_bytes.into()],
    )).await;

    assert!(search_res.is_ok(), "Failed to search in vec_items: {:?}", search_res.err());
    let results = search_res.unwrap();
    assert_eq!(results.len(), 1);

    // Cleanup
    drop(conn);
    // Give a small delay to ensure file handle is released on Windows
    tokio::time::sleep(std::time::Duration::from_millis(500)).await;
    let _ = std::fs::remove_file(test_db);
}