Rust集成测试中无需修改函数签名自动适配数据库连接

方案1：线程本地存储实现连接隔离 #

利用Rust的线程本地存储，让每个测试线程持有独立的数据库连接，生产环境默认加载.env配置，完全不需要修改业务函数签名。

核心实现代码

use std::cell::RefCell;
use sea_orm::{Database, DatabaseConnection, DbErr};

// 线程本地存储，每个线程维护独立的数据库连接
thread_local! {
    static DB_CONN: RefCell<Option<DatabaseConnection>> = RefCell::new(None);
}

/// 初始化数据库连接：生产环境用.env的DATABASE_URL，测试环境传入临时库URL
pub async fn init_db_connection(db_url: Option<&str>) -> Result<(), DbErr> {
    let url = match db_url {
        Some(url) => url.to_string(),
        None => std::env::var("DATABASE_URL").expect("DATABASE_URL未配置"),
    };
    let conn = Database::connect(url).await?;
    DB_CONN.with(|c| *c.borrow_mut() = Some(conn));
    Ok(())
}

/// 业务函数获取连接的统一入口
async fn get_db_conn() -> Result<DatabaseConnection, DbErr> {
    DB_CONN.with(|c| {
        c.borrow().clone()
            .ok_or_else(|| DbErr::Connection("数据库连接未初始化".into()))
    })
}

// 原有业务函数无需修改签名，内部调用get_db_conn即可
pub async fn do_something() -> Result<(), DbErr> {
    let conn = get_db_conn().await?;
    // 业务数据库操作示例：
    // User::find().all(&conn).await?;
    Ok(())
}

pub async fn do_something2() -> Result<(), DbErr> {
    let conn = get_db_conn().await?;
    // 另一项业务逻辑
    Ok(())
}

测试代码示例

#[tokio::test]
async fn test_do_something() {
    // 1. 创建临时PostgreSQL数据库，获取连接URL
    let temp_db_url = create_temp_database().await;
    // 2. 初始化当前测试线程的数据库连接
    init_db_connection(Some(&temp_db_url)).await.unwrap();
    // 3. 直接调用原有业务函数
    do_something().await.unwrap();
    // 4. 可选：清理临时数据库（部分PostgreSQL工具会自动回收）
    drop_temp_database(temp_db_url).await;
}

// 并行测试自动隔离，每个线程拥有独立连接
#[tokio::test]
async fn test_do_something2() {
    let temp_db_url = create_temp_database().await;
    init_db_connection(Some(&temp_db_url)).await.unwrap();
    do_something2().await.unwrap();
}

方案2：全局静态提供者+条件编译 #

通过条件编译区分生产/测试环境，全局维护一个数据库连接提供者，测试时动态替换临时库URL，同样无需修改业务函数签名。

核心实现代码

use once_cell::sync::Lazy;
use std::any::Any;
use std::cell::RefCell;
use sea_orm::{Database, DatabaseConnection, DbErr};

// 定义连接提供者 trait
pub trait DbProvider: Sync + Send {
    async fn get_conn(&self) -> Result<DatabaseConnection, DbErr>;
    fn as_any(&self) -> &dyn Any;
}

// 生产环境提供者：从.env加载配置
struct ProdDbProvider;

impl DbProvider for ProdDbProvider {
    async fn get_conn(&self) -> Result<DatabaseConnection, DbErr> {
        let db_url = std::env::var("DATABASE_URL").expect("DATABASE_URL未配置");
        Database::connect(db_url).await
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

// 测试环境提供者：支持动态设置临时库URL
struct TestDbProvider {
    temp_url: RefCell<Option<String>>,
}

impl TestDbProvider {
    fn new() -> Self {
        Self { temp_url: RefCell::new(None) }
    }

    pub fn set_temp_url(&self, url: String) {
        *self.temp_url.borrow_mut() = Some(url);
    }
}

impl DbProvider for TestDbProvider {
    async fn get_conn(&self) -> Result<DatabaseConnection, DbErr> {
        let url = self.temp_url.borrow().as_ref()
            .expect("测试环境未设置临时数据库URL");
        Database::connect(url).await
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}

// 全局静态提供者：条件编译自动切换生产/测试实现
static DB_PROVIDER: Lazy<Box<dyn DbProvider>> = Lazy::new(|| {
    #[cfg(test)]
    Box::new(TestDbProvider::new())
    #[cfg(not(test))]
    Box::new(ProdDbProvider)
});

// 业务函数直接使用全局提供者
pub async fn do_something() -> Result<(), DbErr> {
    let conn = DB_PROVIDER.get_conn().await?;
    // 业务逻辑
    Ok(())
}

测试代码示例

#[tokio::test]
async fn test_do_something() {
    let temp_db_url = create_temp_database().await;
    // 向下转型为测试提供者，设置临时URL
    if let Some(test_provider) = DB_PROVIDER.as_any().downcast_ref::<TestDbProvider>() {
        test_provider.set_temp_url(temp_db_url);
    }
    do_something().await.unwrap();
}

方案3：临时环境变量隔离（适合简单场景） #

利用temp_env crate隔离每个测试的环境变量，让业务函数从环境变量加载连接时，每个测试拿到不同的临时库URL。注意：此方案不适合缓存连接的业务逻辑，全局缓存的连接会在并行测试中冲突。

测试代码示例

use temp_env::with_var;

#[tokio::test]
async fn test_do_something() {
    let temp_db_url = create_temp_database().await;
    // 临时覆盖当前测试的DATABASE_URL环境变量
    with_var("DATABASE_URL", Some(&temp_db_url), || {
        tokio::runtime::Runtime::new().unwrap().block_on(async {
            do_something().await.unwrap();
        });
    });
}

方案选择建议 #

• 若测试需要并行执行且业务逻辑有连接缓存，优先选择线程本地存储方案，完全隔离每个测试的连接上下文。
• 若需要更灵活的提供者切换，可选择全局静态提供者方案。
• 简单场景且无连接缓存时，可使用临时环境变量方案。

内容的提问来源于stack exchange，提问作者spacecodeur