tests.rs

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