codex-memory 3.0.15

use anyhow::Result;
use sqlx::postgres::{PgPool, PgPoolOptions};
use sqlx::Row;
use std::env;
use std::sync::Once;
use uuid::Uuid;

static INIT: Once = Once::new();

/// Ensure .env is loaded for tests
fn ensure_env_loaded() {
    INIT.call_once(|| {
        dotenvy::dotenv().ok();
    });
}

/// Test database manager for isolated test execution
/// Creates test-specific tables with UUID prefixes to ensure complete isolation
pub struct TestDatabaseManager {
    pool: Option<PgPool>,
    pub test_id: String,
    pub memories_table: String,
}

impl TestDatabaseManager {
    /// Create a new test database manager with isolated tables
    pub fn new() -> Result<Self> {
        // Ensure .env is loaded
        ensure_env_loaded();

        // Generate unique test ID for this test run
        let test_id = Uuid::new_v4().to_string().replace("-", "_");
        let memories_table = format!("test_memories_{}", test_id);

        Ok(Self {
            pool: None,
            test_id,
            memories_table,
        })
    }

    /// Get the test database URL
    pub fn get_test_database_url(&self) -> String {
        // Ensure .env is loaded
        ensure_env_loaded();

        env::var("TEST_DATABASE_URL")
            .expect("TEST_DATABASE_URL must be set - NEVER use production database for tests! Set TEST_DATABASE_URL=postgresql://user:pass@host:port/codex_test_db")
    }

    /// Setup test database with isolated test tables
    pub async fn setup_test_database(&mut self) -> Result<PgPool> {
        // CRITICAL: Always use TEST_DATABASE_URL - NEVER fall back to production
        let test_url = self.get_test_database_url();

        // Create test-specific schema to isolate test data
        let schema_name = format!("test_schema_{}", self.test_id);

        // First, create a temporary pool to set up schema
        let temp_pool = PgPoolOptions::new()
            .max_connections(1)
            .connect(&test_url)
            .await?;

        // Drop existing schema if it exists (for clean test run)
        sqlx::query(&format!("DROP SCHEMA IF EXISTS {} CASCADE", schema_name))
            .execute(&temp_pool)
            .await
            .ok();

        // Create isolated schema for this test
        sqlx::query(&format!("CREATE SCHEMA {}", schema_name))
            .execute(&temp_pool)
            .await?;

        // Close temporary pool
        temp_pool.close().await;

        // Now create the main pool with search_path set for ALL connections
        let search_path_setting = format!("SET search_path TO {}, public", schema_name);
        let pool = PgPoolOptions::new()
            .max_connections(5)
            .after_connect(move |conn, _meta| {
                let search_path_cmd = search_path_setting.clone();
                Box::pin(async move {
                    sqlx::query(&search_path_cmd).execute(conn).await?;
                    Ok(())
                })
            })
            .connect(&test_url)
            .await?;

        // Run actual migrations in the test schema to ensure proper structure
        self.run_migrations(&pool).await?;

        self.pool = Some(pool.clone());
        Ok(pool)
    }

    /// Run migrations for test database
    async fn run_migrations(&self, pool: &PgPool) -> Result<()> {
        // Run the actual production migrations to ensure test schema matches production
        codex_memory::database::run_migrations(pool).await?;
        Ok(())
    }

    /// Clean up test database/tables
    pub async fn cleanup(&mut self) -> Result<()> {
        if let Some(pool) = &self.pool {
            // Drop the test schema cascade (removes all tables in it)
            let schema_name = format!("test_schema_{}", self.test_id);
            let drop_schema = format!("DROP SCHEMA IF EXISTS {} CASCADE", schema_name);

            sqlx::query(&drop_schema).execute(pool).await.ok(); // Ignore errors during cleanup

            println!("Cleaned up test schema: {}", schema_name);
        }

        Ok(())
    }
}

impl Drop for TestDatabaseManager {
    fn drop(&mut self) {
        // Close the pool if it exists to prevent hanging connections
        if let Some(pool) = &self.pool {
            // Note: close() returns a future but we can't await in Drop
            // The connection pool will be closed when the last reference is dropped
            let _close_future = pool.close();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_database_manager_lifecycle() -> Result<()> {
        let mut manager = TestDatabaseManager::new()?;

        // Verify unique test ID generation
        assert!(!manager.test_id.is_empty());
        assert!(!manager.memories_table.is_empty());

        // Setup test database
        let pool = manager.setup_test_database().await?;

        // Debug: Check current search path
        let search_path_result = sqlx::query("SHOW search_path").fetch_one(&pool).await?;
        let search_path: String = search_path_result.try_get("search_path")?;
        println!("Current search_path: {}", search_path);

        // Debug: Check if memories table exists in test schema
        let test_schema = format!("test_schema_{}", manager.test_id);
        let table_check = sqlx::query(&format!(
            "SELECT COUNT(*) as count FROM information_schema.tables 
             WHERE table_schema = '{}' AND table_name = 'memories'",
            test_schema
        ))
        .fetch_one(&pool)
        .await?;
        let table_exists: i64 = table_check.try_get("count")?;
        println!(
            "Memories table exists in {}: {}",
            test_schema,
            table_exists > 0
        );

        // Verify we can query the test schema
        let result = sqlx::query(&format!(
            "SELECT COUNT(*) as count FROM {}.memories",
            test_schema
        ))
        .fetch_one(&pool)
        .await?;

        // Should have zero rows in fresh test table
        let count: i64 = result.try_get("count")?;
        assert_eq!(count, 0);

        // Insert test data with unique hash based on test ID
        let unique_content = format!("Test content {}", manager.test_id);
        let unique_hash = format!("test_hash_{}", manager.test_id);
        sqlx::query(&format!(
            "INSERT INTO {}.memories (content, content_hash, context, summary) VALUES ($1, $2, $3, $4)",
            test_schema
        ))
        .bind(&unique_content)
        .bind(&unique_hash)
        .bind("Test context")
        .bind("Test summary")
        .execute(&pool)
        .await?;

        // Verify insertion
        let result = sqlx::query(&format!(
            "SELECT COUNT(*) as count FROM {}.memories",
            test_schema
        ))
        .fetch_one(&pool)
        .await?;
        let count: i64 = result.try_get("count")?;
        assert_eq!(count, 1);

        // Cleanup
        manager.cleanup().await?;

        // After cleanup, the schema should be gone
        // (Can't easily test this without reconnecting)

        Ok(())
    }
}