//! Test Database Utilities for Foxhunt HFT Trading System //! //! This module provides database setup, cleanup, and helper utilities //! for testing database operations in isolation. use std::sync::Arc; use sqlx::{PgPool, Postgres, migrate::MigrateDatabase}; use tokio::sync::OnceCell; use uuid::Uuid; use super::test_config::TestConfig; /// Test database manager for setup and cleanup #[derive(Debug)] pub struct TestDatabase { pub pool: PgPool, pub config: TestConfig, pub database_name: String, cleanup_on_drop: bool, } impl TestDatabase { /// Create a new test database with unique name pub async fn new() -> Result { Self::with_config(TestConfig::for_unit_tests()).await } /// Create a test database with specific configuration pub async fn with_config(config: TestConfig) -> Result { let database_name = format!("test_db_{}", Uuid::new_v4().simple()); let base_url = Self::get_base_database_url(&config.database.url); let full_url = format!("{}/{}", base_url, database_name); // Create the test database Postgres::create_database(&full_url).await?; // Connect to the new database let pool = PgPool::connect(&full_url).await?; // Run migrations if enabled if config.database.auto_migrate { Self::run_migrations(&pool).await?; } Ok(Self { pool, config, database_name, cleanup_on_drop: true, }) } /// Create a test database from existing pool (for shared testing) pub fn from_pool(pool: PgPool, config: TestConfig) -> Self { Self { pool, config, database_name: "shared_test_db".to_string(), cleanup_on_drop: false, // Don't cleanup shared pools } } /// Get the database pool pub fn pool(&self) -> &PgPool { &self.pool } /// Execute SQL and return affected rows pub async fn execute(&self, sql: &str) -> Result { let result = sqlx::query(sql).execute(&self.pool).await?; Ok(result.rows_affected()) } /// Fetch one row pub async fn fetch_one(&self, sql: &str) -> Result where T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin, { sqlx::query_as::<_, T>(sql).fetch_one(&self.pool).await } /// Fetch all rows pub async fn fetch_all(&self, sql: &str) -> Result, sqlx::Error> where T: for<'r> sqlx::FromRow<'r, sqlx::postgres::PgRow> + Send + Unpin, { sqlx::query_as::<_, T>(sql).fetch_all(&self.pool).await } /// Clean all test data from database (without dropping) pub async fn clean_test_data(&self) -> Result<(), sqlx::Error> { // Clean in dependency order (children first, parents last) let cleanup_queries = vec![ "DELETE FROM stress_test_results WHERE portfolio_id LIKE 'TEST_%'", "DELETE FROM counterparty_exposure WHERE counterparty_id LIKE 'TEST_%'", "DELETE FROM risk_factor_exposure WHERE portfolio_id LIKE 'TEST_%'", "DELETE FROM portfolio_performance WHERE portfolio_id LIKE 'TEST_%'", "DELETE FROM positions WHERE portfolio_id LIKE 'TEST_%'", "DELETE FROM portfolios WHERE id LIKE 'TEST_%'", "DELETE FROM counterparty WHERE id LIKE 'TEST_%'", "DELETE FROM instruments WHERE symbol LIKE 'TEST_%'", "DELETE FROM stress_scenarios WHERE name LIKE 'Test %'", "DELETE FROM risk_calculation_jobs WHERE created_by = 'test_system'", "DELETE FROM market_data_feeds WHERE provider_name LIKE 'TEST_%'", "DELETE FROM liquidity_metrics WHERE symbol LIKE 'TEST_%'", "DELETE FROM economic_scenarios WHERE created_by = 'test_system'", "DELETE FROM risk_reports WHERE generated_by = 'test_system'", "DELETE FROM risk_report_templates WHERE created_by = 'test_system'", "DELETE FROM custom_risk_metric_results WHERE portfolio_id LIKE 'TEST_%'", "DELETE FROM custom_risk_metrics WHERE created_by = 'test_system'", ]; for query in cleanup_queries { self.execute(query).await?; } Ok(()) } /// Insert test data into database pub async fn insert_test_data(&self) -> Result<(), sqlx::Error> { // Insert test instruments self.insert_test_instruments().await?; // Insert test portfolios self.insert_test_portfolios().await?; // Insert test counterparties self.insert_test_counterparties().await?; Ok(()) } /// Insert test instruments async fn insert_test_instruments(&self) -> Result<(), sqlx::Error> { use super::{ALL_TEST_SYMBOLS, builders::BatchBuilder}; let instruments = BatchBuilder::create_diverse_instruments(ALL_TEST_SYMBOLS.len()); for instrument in instruments { sqlx::query( r#" INSERT INTO instruments ( id, symbol, name, instrument_type, asset_class, currency, is_active, created_at, updated_at, metadata ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) ON CONFLICT (symbol) DO NOTHING "#, ) .bind(instrument.id) .bind(instrument.symbol) .bind(instrument.name) .bind(instrument.instrument_type.to_string()) .bind(instrument.asset_class.to_string()) .bind(instrument.currency) .bind(instrument.is_active) .bind(instrument.created_at) .bind(instrument.updated_at) .bind(instrument.metadata) .execute(&self.pool) .await?; } Ok(()) } /// Insert test portfolios async fn insert_test_portfolios(&self) -> Result<(), sqlx::Error> { use super::{TEST_PORTFOLIO_1, TEST_PORTFOLIO_2, builders::PortfolioBuilder}; let portfolios = vec![ PortfolioBuilder::new().with_id(TEST_PORTFOLIO_1).build(), PortfolioBuilder::new().with_id(TEST_PORTFOLIO_2).build(), ]; for portfolio in portfolios { sqlx::query( r#" INSERT INTO portfolios ( id, name, description, base_currency, portfolio_type, inception_date, manager_id, is_active, created_at, updated_at, metadata ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) ON CONFLICT (id) DO NOTHING "#, ) .bind(portfolio.id) .bind(portfolio.name) .bind(portfolio.description) .bind(portfolio.base_currency) .bind(portfolio.portfolio_type) .bind(portfolio.inception_date) .bind(portfolio.manager_id) .bind(portfolio.is_active) .bind(portfolio.created_at) .bind(portfolio.updated_at) .bind(portfolio.metadata) .execute(&self.pool) .await?; } Ok(()) } /// Insert test counterparties async fn insert_test_counterparties(&self) -> Result<(), sqlx::Error> { use super::builders::CounterpartyBuilder; let counterparties = vec![ CounterpartyBuilder::new().with_id("TEST_COUNTERPARTY_001").bank().build(), CounterpartyBuilder::new().with_id("TEST_COUNTERPARTY_002").broker().build(), CounterpartyBuilder::new().with_id("TEST_COUNTERPARTY_003").exchange().build(), ]; for counterparty in counterparties { sqlx::query( r#" INSERT INTO counterparty ( id, name, counterparty_type, country, credit_rating, is_active, netting_agreement, created_at, updated_at, metadata ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10) ON CONFLICT (id) DO NOTHING "#, ) .bind(counterparty.id) .bind(counterparty.name) .bind(counterparty.counterparty_type) .bind(counterparty.country) .bind(counterparty.credit_rating) .bind(counterparty.is_active) .bind(counterparty.netting_agreement) .bind(counterparty.created_at) .bind(counterparty.updated_at) .bind(counterparty.metadata) .execute(&self.pool) .await?; } Ok(()) } /// Check if database schema is ready pub async fn is_schema_ready(&self) -> bool { let result: Result<(bool,), _> = sqlx::query_as( "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = 'instruments')" ) .fetch_one(&self.pool) .await; match result { Ok((exists,)) => exists, Err(_) => false, } } /// Get database statistics pub async fn get_stats(&self) -> Result { let instruments_count: Option = sqlx::query_scalar( "SELECT COUNT(*) FROM instruments WHERE symbol LIKE 'TEST_%'" ) .fetch_one(&self.pool) .await?; let portfolios_count: Option = sqlx::query_scalar( "SELECT COUNT(*) FROM portfolios WHERE id LIKE 'TEST_%'" ) .fetch_one(&self.pool) .await?; let positions_count: Option = sqlx::query_scalar( "SELECT COUNT(*) FROM positions WHERE portfolio_id LIKE 'TEST_%'" ) .fetch_one(&self.pool) .await?; Ok(DatabaseStats { instruments_count: instruments_count.unwrap_or(0), portfolios_count: portfolios_count.unwrap_or(0), positions_count: positions_count.unwrap_or(0), }) } /// Run database migrations async fn run_migrations(pool: &PgPool) -> Result<(), sqlx::Error> { // Basic schema creation (simplified for testing) sqlx::query( r#" CREATE TABLE IF NOT EXISTS instruments ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), symbol VARCHAR(50) UNIQUE NOT NULL, name VARCHAR(255) NOT NULL, instrument_type VARCHAR(50) NOT NULL, asset_class VARCHAR(50) NOT NULL, currency CHAR(3) NOT NULL, is_active BOOLEAN DEFAULT true, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW(), metadata JSONB DEFAULT '{}' ) "# ) .execute(pool) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS portfolios ( id VARCHAR(50) PRIMARY KEY, name VARCHAR(255) NOT NULL, description TEXT, base_currency CHAR(3) NOT NULL, portfolio_type VARCHAR(100) NOT NULL, inception_date TIMESTAMPTZ NOT NULL, manager_id VARCHAR(100) NOT NULL, is_active BOOLEAN DEFAULT true, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW(), metadata JSONB DEFAULT '{}' ) "# ) .execute(pool) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS positions ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), portfolio_id VARCHAR(50) REFERENCES portfolios(id), symbol VARCHAR(50) NOT NULL, quantity DECIMAL NOT NULL, average_price DECIMAL NOT NULL, market_price DECIMAL NOT NULL, market_value DECIMAL NOT NULL, unrealized_pnl DECIMAL NOT NULL, currency CHAR(3) NOT NULL, entry_date TIMESTAMPTZ NOT NULL, last_updated TIMESTAMPTZ DEFAULT NOW() ) "# ) .execute(pool) .await?; sqlx::query( r#" CREATE TABLE IF NOT EXISTS counterparty ( id VARCHAR(50) PRIMARY KEY, name VARCHAR(255) NOT NULL, counterparty_type VARCHAR(100) NOT NULL, country CHAR(2) NOT NULL, credit_rating VARCHAR(10), is_active BOOLEAN DEFAULT true, netting_agreement BOOLEAN DEFAULT false, created_at TIMESTAMPTZ DEFAULT NOW(), updated_at TIMESTAMPTZ DEFAULT NOW(), metadata JSONB DEFAULT '{}' ) "# ) .execute(pool) .await?; Ok(()) } /// Get base database URL without database name fn get_base_database_url(full_url: &str) -> String { if let Some(last_slash) = full_url.rfind('/') { full_url[..last_slash].to_string() } else { full_url.to_string() } } } impl Drop for TestDatabase { fn drop(&mut self) { if self.cleanup_on_drop && self.database_name != "shared_test_db" { // Note: We can't use async in Drop, so we schedule cleanup // In practice, you might want to use a cleanup service or manual cleanup eprintln!("TestDatabase: Consider cleaning up database: {}", self.database_name); } } } /// Database statistics for monitoring test data #[derive(Debug, Clone)] pub struct DatabaseStats { pub instruments_count: i64, pub portfolios_count: i64, pub positions_count: i64, } /// Shared test database pool for tests that don't need isolation static SHARED_TEST_POOL: OnceCell> = OnceCell::const_new(); /// Get or create shared test database pool pub async fn get_shared_test_db() -> Result, sqlx::Error> { SHARED_TEST_POOL .get_or_try_init(|| async { let config = TestConfig::for_integration_tests(); let test_db = TestDatabase::with_config(config).await?; test_db.insert_test_data().await?; Ok(Arc::new(test_db)) }) .await .map(Arc::clone) } /// Test database transaction for isolated test operations #[derive(Debug)] pub struct TestTransaction<'a> { tx: sqlx::Transaction<'a, Postgres>, } impl<'a> TestTransaction<'a> { /// Begin a new test transaction pub async fn begin(pool: &'a PgPool) -> Result { let tx = pool.begin().await?; Ok(Self { tx }) } /// Execute SQL within transaction pub async fn execute(&mut self, sql: &str) -> Result { let result = sqlx::query(sql).execute(&mut *self.tx).await?; Ok(result.rows_affected()) } /// Commit the transaction pub async fn commit(self) -> Result<(), sqlx::Error> { self.tx.commit().await } /// Rollback the transaction pub async fn rollback(self) -> Result<(), sqlx::Error> { self.tx.rollback().await } } /// Utility macro for test database setup #[macro_export] macro_rules! setup_test_db { () => {{ let test_db = $crate::fixtures::test_database::TestDatabase::new().await?; test_db.insert_test_data().await?; test_db }}; ($config:expr) => {{ let test_db = $crate::fixtures::test_database::TestDatabase::with_config($config).await?; test_db.insert_test_data().await?; test_db }}; } /// Utility macro for test transaction #[macro_export] macro_rules! test_transaction { ($db:expr, $code:block) => {{ let mut tx = $crate::fixtures::test_database::TestTransaction::begin($db.pool()).await?; let result = async move { $code }.await; tx.rollback().await?; // Always rollback test transactions result }}; } #[cfg(test)] mod tests { use super::*; /// Helper to check if PostgreSQL is available async fn is_postgres_available() -> bool { let config = TestConfig::for_unit_tests(); PgPool::connect(&config.database.url).await.is_ok() } #[tokio::test] #[ignore] async fn test_database_creation() { if !is_postgres_available().await { eprintln!("Skipping test: PostgreSQL not available"); return; } let test_db = TestDatabase::new().await.unwrap(); assert!(test_db.is_schema_ready().await); assert!(test_db.database_name.starts_with("test_db_")); } #[tokio::test] #[ignore] async fn test_insert_and_clean_test_data() { if !is_postgres_available().await { eprintln!("Skipping test: PostgreSQL not available"); return; } let test_db = TestDatabase::new().await.unwrap(); // Insert test data test_db.insert_test_data().await.unwrap(); // Check data was inserted let stats = test_db.get_stats().await.unwrap(); assert!(stats.instruments_count > 0); assert!(stats.portfolios_count > 0); // Clean test data test_db.clean_test_data().await.unwrap(); // Check data was cleaned let stats_after = test_db.get_stats().await.unwrap(); assert_eq!(stats_after.instruments_count, 0); assert_eq!(stats_after.portfolios_count, 0); } #[tokio::test] #[ignore] async fn test_database_stats() { if !is_postgres_available().await { eprintln!("Skipping test: PostgreSQL not available"); return; } let test_db = TestDatabase::new().await.unwrap(); test_db.insert_test_data().await.unwrap(); let stats = test_db.get_stats().await.unwrap(); assert!(stats.instruments_count >= 0); assert!(stats.portfolios_count >= 0); assert!(stats.positions_count >= 0); } #[tokio::test] #[ignore] async fn test_transaction_rollback() { if !is_postgres_available().await { eprintln!("Skipping test: PostgreSQL not available"); return; } let test_db = TestDatabase::new().await.unwrap(); // Count before transaction let initial_count = test_db.execute("SELECT COUNT(*) FROM instruments").await.unwrap(); // Start transaction and insert data let mut tx = TestTransaction::begin(test_db.pool()).await.unwrap(); tx.execute("INSERT INTO instruments (symbol, name, instrument_type, asset_class, currency) VALUES ('TX_TEST', 'Transaction Test', 'equity', 'equities', 'USD')").await.unwrap(); tx.rollback().await.unwrap(); // Count after rollback should be the same let final_count = test_db.execute("SELECT COUNT(*) FROM instruments").await.unwrap(); assert_eq!(initial_count, final_count); } }