minne/common/src/test_utils.rs

//! Shared helpers for in-memory SurrealDB tests.
#![cfg(any(test, feature = "test-utils"))]

use anyhow::{Context, Result};
use uuid::Uuid;

use crate::storage::{
    db::SurrealDbClient,
    indexes::{ensure_runtime, rebuild},
    types::{
        knowledge_entity_embedding::KnowledgeEntityEmbedding, system_settings::SystemSettings,
        text_chunk_embedding::TextChunkEmbedding,
    },
};

const TEST_NAMESPACE: &str = "test_ns";

/// Starts an in-memory database, applies migrations, and returns a client.
///
/// # Errors
///
/// Returns an error if the database cannot be started or migrations fail.
pub async fn setup_test_db() -> Result<SurrealDbClient> {
    let database = Uuid::new_v4().to_string();
    let db = SurrealDbClient::memory(TEST_NAMESPACE, &database)
        .await
        .context("start in-memory surrealdb")?;

    db.apply_migrations().await.context("apply migrations")?;

    Ok(db)
}

/// Updates singleton [`SystemSettings`] embedding dimensions for tests.
///
/// # Errors
///
/// Returns an error if settings cannot be loaded or updated.
pub async fn configure_embedding_dimension(db: &SurrealDbClient, dimension: u32) -> Result<()> {
    let mut settings = SystemSettings::get_current(db).await?;
    settings.embedding_dimensions = dimension;
    SystemSettings::update(db, settings).await?;
    Ok(())
}

/// Starts a test database and sets the embedding dimension in system settings.
///
/// # Errors
///
/// Returns an error if setup or settings update fails.
pub async fn setup_test_db_with_embedding_dimension(dimension: u32) -> Result<SurrealDbClient> {
    let db = setup_test_db().await?;
    configure_embedding_dimension(&db, dimension).await?;
    Ok(db)
}

/// Prepares a database for text-chunk embedding tests at the given dimension.
///
/// # Errors
///
/// Returns an error if setup, settings update, or index redefinition fails.
pub async fn prepare_text_chunk_test_db(dimension: u32) -> Result<SurrealDbClient> {
    let db = setup_test_db_with_embedding_dimension(dimension).await?;
    TextChunkEmbedding::redefine_hnsw_index(&db, dimension as usize)
        .await
        .with_context(|| format!("set text chunk index dimension to {dimension}"))?;
    Ok(db)
}

/// Prepares a database for knowledge-entity embedding tests at the given dimension.
///
/// # Errors
///
/// Returns an error if setup, settings update, or index redefinition fails.
pub async fn prepare_knowledge_entity_test_db(dimension: u32) -> Result<SurrealDbClient> {
    let db = setup_test_db_with_embedding_dimension(dimension).await?;
    KnowledgeEntityEmbedding::redefine_hnsw_index(&db, dimension as usize)
        .await
        .with_context(|| format!("set knowledge entity index dimension to {dimension}"))?;
    Ok(db)
}

/// Starts a test database and ensures runtime FTS/HNSW indexes are ready.
///
/// # Errors
///
/// Returns an error if setup, index creation, or rebuild fails.
pub async fn setup_test_db_with_runtime_indexes() -> Result<SurrealDbClient> {
    let db = setup_test_db().await?;
    ensure_runtime(&db, 1536).await?;
    rebuild(&db).await?;
    Ok(db)
}