Set up a test database once and get connections for tests that run in a diesel test_transaction

TestDatabase.rs

use std::error::Error;
use diesel::{sql_query, Connection, RunQueryDsl, insert_into};
use diesel::{PgConnection};
use diesel::pg::Pg;
use diesel::r2d2::ConnectionManager;
use diesel_migrations::{EmbeddedMigrations, MigrationHarness};
use r2d2::CustomizeConnection;
use reqwest::Url;
use lazy_static::lazy_static;
use log::{info};
use std::time::{Instant};

pub type Pool = r2d2::Pool<ConnectionManager<PgConnection>>;

fn run_migrations(
    connection: &mut impl MigrationHarness<Pg>
) -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
    use diesel_migrations::{embed_migrations};
    const MIGRATIONS: EmbeddedMigrations = embed_migrations!();
    connection.run_pending_migrations(MIGRATIONS)?;

    Ok(())
}

fn create_initial_data(conn: &mut PgConnection) -> Result<(), Box<dyn Error>> {
    use crate::schema::member_member::dsl::*;
    use crate::models::member::{MemberMember, MemberMemberForm, MemberTenantForm, StaffuserForm};

    let json = r#"{
        "companycode": "ritehite",
        "is_deleted": false,
        "member_type": "maintenance"
    }"#;
    info!("creating ritehite member");
    let member_form_ritehite = serde_json::from_str::<MemberMemberForm>(json)?;
    let ritehite_member = insert_into(member_member)
        .values(&member_form_ritehite)
        .get_result::<MemberMember>(conn)?;
    info!("creating ritehite member done");

    info!("creating ritehite tenant");
    MemberTenantForm::create("ritehite", ritehite_member.id, conn)?;
    info!("creating ritehite tenant done");

    info!("creating staff user");
    StaffuserForm::create(
        conn, "staff_user", "test", "Staff", "User",
        "[email protected]", false, true, true
    )?;
    info!("creating staff user done");

    Ok(())
}

#[derive(Debug, Clone)]
pub struct TestDatabase {
    url: Url
}

impl TestDatabase {
    pub fn get() -> Self {
        let start = Instant::now();
        dotenv::dotenv().ok();
        let name = "my24_live_test";
        let default_db_url = std::env::var("DATABASE_URL").expect("Failed to get DATABASE_URL from env");
        let mut conn = PgConnection::establish(&default_db_url).unwrap();

        // drop test db if exists
        info!("drop test db if exists");
        sql_query(format!("DROP DATABASE IF EXISTS {};", &name))
            .execute(&mut conn)
            .unwrap();

        // create test db
        info!("create test db");
        sql_query(format!("CREATE DATABASE {};", &name))
            .execute(&mut conn)
            .unwrap();

        let mut url = Url::parse(&default_db_url).unwrap();
        url.set_path(&name);

        info!("connect to test db");
        let mut conn = PgConnection::establish(&(url.to_string().clone())).unwrap();

        // run migrations
        info!("run migrations");
        run_migrations(&mut conn).expect("Unable to run migrations");

        // create test tenants
        info!("fill test db");
        create_initial_data(&mut conn).expect("Unable to create test tenants");

        let duration = start.elapsed();
        info!("Creating test database, schemas and fixtures done, took {:?}!", duration);

        Self {
            url
        }
    }

    pub fn get_pool(self) -> Pool {
        let manager = ConnectionManager::<PgConnection>::new(self.url.to_string());
        let mut builder = r2d2::Pool::builder();

        builder = builder.connection_customizer(Box::new(TestConnectionCustomizer));
        builder
            .max_size(1)
            .build(manager)
            .expect("Failed to create database connection pool")
    }

}

lazy_static! {
    pub static ref TESTDB: TestDatabase = {
        TestDatabase::get()
    };
}

#[derive(Debug, Clone, Copy)]
pub struct TestConnectionCustomizer;

impl<C, E> CustomizeConnection<C, E> for TestConnectionCustomizer
    where
        C: diesel::Connection,
{
    fn on_acquire(&self, conn: &mut C) -> Result<(), E> {
        conn.begin_test_transaction()
            .expect("Failed to start test transaction");

        Ok(())
    }
}

      
# Usage
let test_db = TESTDB.clone();
let pool = test_db.get_pool();