//! 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 { 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 { 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 { 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 { 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 { let db = setup_test_db().await?; ensure_runtime(&db, 1536).await?; rebuild(&db).await?; Ok(db) }