tests.rs

  1mod buffer_tests;
  2mod db_tests;
  3mod feature_flag_tests;
  4mod message_tests;
  5
  6use super::*;
  7use gpui::executor::Background;
  8use parking_lot::Mutex;
  9use sea_orm::ConnectionTrait;
 10use sqlx::migrate::MigrateDatabase;
 11use std::sync::Arc;
 12
 13pub struct TestDb {
 14    pub db: Option<Arc<Database>>,
 15    pub connection: Option<sqlx::AnyConnection>,
 16}
 17
 18impl TestDb {
 19    pub fn sqlite(background: Arc<Background>) -> Self {
 20        let url = format!("sqlite::memory:");
 21        let runtime = tokio::runtime::Builder::new_current_thread()
 22            .enable_io()
 23            .enable_time()
 24            .build()
 25            .unwrap();
 26
 27        let mut db = runtime.block_on(async {
 28            let mut options = ConnectOptions::new(url);
 29            options.max_connections(5);
 30            let db = Database::new(options, Executor::Deterministic(background))
 31                .await
 32                .unwrap();
 33            let sql = include_str!(concat!(
 34                env!("CARGO_MANIFEST_DIR"),
 35                "/migrations.sqlite/20221109000000_test_schema.sql"
 36            ));
 37            db.pool
 38                .execute(sea_orm::Statement::from_string(
 39                    db.pool.get_database_backend(),
 40                    sql.into(),
 41                ))
 42                .await
 43                .unwrap();
 44            db
 45        });
 46
 47        db.runtime = Some(runtime);
 48
 49        Self {
 50            db: Some(Arc::new(db)),
 51            connection: None,
 52        }
 53    }
 54
 55    pub fn postgres(background: Arc<Background>) -> Self {
 56        static LOCK: Mutex<()> = Mutex::new(());
 57
 58        let _guard = LOCK.lock();
 59        let mut rng = StdRng::from_entropy();
 60        let url = format!(
 61            "postgres://postgres@localhost/zed-test-{}",
 62            rng.gen::<u128>()
 63        );
 64        let runtime = tokio::runtime::Builder::new_current_thread()
 65            .enable_io()
 66            .enable_time()
 67            .build()
 68            .unwrap();
 69
 70        let mut db = runtime.block_on(async {
 71            sqlx::Postgres::create_database(&url)
 72                .await
 73                .expect("failed to create test db");
 74            let mut options = ConnectOptions::new(url);
 75            options
 76                .max_connections(5)
 77                .idle_timeout(Duration::from_secs(0));
 78            let db = Database::new(options, Executor::Deterministic(background))
 79                .await
 80                .unwrap();
 81            let migrations_path = concat!(env!("CARGO_MANIFEST_DIR"), "/migrations");
 82            db.migrate(Path::new(migrations_path), false).await.unwrap();
 83            db
 84        });
 85
 86        db.runtime = Some(runtime);
 87
 88        Self {
 89            db: Some(Arc::new(db)),
 90            connection: None,
 91        }
 92    }
 93
 94    pub fn db(&self) -> &Arc<Database> {
 95        self.db.as_ref().unwrap()
 96    }
 97}
 98
 99#[macro_export]
100macro_rules! test_both_dbs {
101    ($test_name:ident, $postgres_test_name:ident, $sqlite_test_name:ident) => {
102        #[gpui::test]
103        async fn $postgres_test_name() {
104            let test_db = crate::db::TestDb::postgres(
105                gpui::executor::Deterministic::new(0).build_background(),
106            );
107            $test_name(test_db.db()).await;
108        }
109
110        #[gpui::test]
111        async fn $sqlite_test_name() {
112            let test_db =
113                crate::db::TestDb::sqlite(gpui::executor::Deterministic::new(0).build_background());
114            $test_name(test_db.db()).await;
115        }
116    };
117}
118
119impl Drop for TestDb {
120    fn drop(&mut self) {
121        let db = self.db.take().unwrap();
122        if let sea_orm::DatabaseBackend::Postgres = db.pool.get_database_backend() {
123            db.runtime.as_ref().unwrap().block_on(async {
124                use util::ResultExt;
125                let query = "
126                        SELECT pg_terminate_backend(pg_stat_activity.pid)
127                        FROM pg_stat_activity
128                        WHERE
129                            pg_stat_activity.datname = current_database() AND
130                            pid <> pg_backend_pid();
131                    ";
132                db.pool
133                    .execute(sea_orm::Statement::from_string(
134                        db.pool.get_database_backend(),
135                        query.into(),
136                    ))
137                    .await
138                    .log_err();
139                sqlx::Postgres::drop_database(db.options.get_url())
140                    .await
141                    .log_err();
142            })
143        }
144    }
145}