From e3f055d9501ebd0cb703fae17366629436d3072c Mon Sep 17 00:00:00 2001 From: Max Brunsfeld Date: Mon, 7 Feb 2022 15:00:00 -0800 Subject: [PATCH] Use a pool of databases to speed up integration tests Also, use env_logger consistently in the tests for each crate. Only initiallize the logger at all if some RUST_LOG env var is set. Co-Authored-By: Nathan Sobo --- crates/editor/src/test.rs | 5 +- crates/gpui/src/test.rs | 6 +- crates/language/src/tests.rs | 5 +- crates/server/src/db.rs | 116 ++++++++++++++++++++++++----------- crates/server/src/rpc.rs | 5 +- crates/text/src/tests.rs | 5 +- crates/zed/src/test.rs | 4 +- 7 files changed, 99 insertions(+), 47 deletions(-) diff --git a/crates/editor/src/test.rs b/crates/editor/src/test.rs index f4622d1f6e..67622db83f 100644 --- a/crates/editor/src/test.rs +++ b/crates/editor/src/test.rs @@ -1,6 +1,7 @@ #[cfg(test)] #[ctor::ctor] fn init_logger() { - // std::env::set_var("RUST_LOG", "info"); - env_logger::init(); + if std::env::var("RUST_LOG").is_ok() { + env_logger::init(); + } } diff --git a/crates/gpui/src/test.rs b/crates/gpui/src/test.rs index b4b4e621ac..af6430d36c 100644 --- a/crates/gpui/src/test.rs +++ b/crates/gpui/src/test.rs @@ -18,9 +18,9 @@ use crate::{ #[cfg(test)] #[ctor::ctor] fn init_logger() { - env_logger::builder() - .filter_level(log::LevelFilter::Info) - .init(); + if std::env::var("RUST_LOG").is_ok() { + env_logger::init(); + } } pub fn run_test( diff --git a/crates/language/src/tests.rs b/crates/language/src/tests.rs index 1fe3c7a917..40dfd06e94 100644 --- a/crates/language/src/tests.rs +++ b/crates/language/src/tests.rs @@ -17,8 +17,9 @@ use util::test::Network; #[cfg(test)] #[ctor::ctor] fn init_logger() { - // std::env::set_var("RUST_LOG", "info"); - env_logger::init(); + if std::env::var("RUST_LOG").is_ok() { + env_logger::init(); + } } #[test] diff --git a/crates/server/src/db.rs b/crates/server/src/db.rs index f71f40efd0..b50d7daf46 100644 --- a/crates/server/src/db.rs +++ b/crates/server/src/db.rs @@ -526,54 +526,88 @@ pub struct ChannelMessage { #[cfg(test)] pub mod tests { use super::*; + use lazy_static::lazy_static; + use parking_lot::Mutex; use rand::prelude::*; use sqlx::{ migrate::{MigrateDatabase, Migrator}, Postgres, }; - use std::path::Path; + use std::{mem, path::Path}; pub struct TestDb { - pub db: Db, + pub db: Option, pub name: String, pub url: String, } - impl TestDb { - pub fn new() -> Self { - // Enable tests to run in parallel by serializing the creation of each test database. - lazy_static::lazy_static! { - static ref DB_CREATION: std::sync::Mutex<()> = std::sync::Mutex::new(()); - } + lazy_static! { + static ref POOL: Mutex> = Default::default(); + } - let mut rng = StdRng::from_entropy(); - let name = format!("zed-test-{}", rng.gen::()); - let url = format!("postgres://postgres@localhost/{}", name); - let migrations_path = Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations")); - let db = block_on(async { - { - let _lock = DB_CREATION.lock(); - Postgres::create_database(&url) - .await - .expect("failed to create test db"); - } - let mut db = Db::new(&url, 5).await.unwrap(); - db.test_mode = true; - let migrator = Migrator::new(migrations_path).await.unwrap(); - migrator.run(&db.pool).await.unwrap(); - db - }); - - Self { db, name, url } - } - - pub fn db(&self) -> &Db { - &self.db + #[ctor::dtor] + fn clear_pool() { + for db in POOL.lock().drain(..) { + db.teardown(); } } - impl Drop for TestDb { - fn drop(&mut self) { + impl TestDb { + pub fn new() -> Self { + let mut pool = POOL.lock(); + if let Some(db) = pool.pop() { + db.truncate(); + db + } else { + let mut rng = StdRng::from_entropy(); + let name = format!("zed-test-{}", rng.gen::()); + let url = format!("postgres://postgres@localhost/{}", name); + let migrations_path = Path::new(concat!(env!("CARGO_MANIFEST_DIR"), "/migrations")); + let db = block_on(async { + Postgres::create_database(&url) + .await + .expect("failed to create test db"); + let mut db = Db::new(&url, 5).await.unwrap(); + db.test_mode = true; + let migrator = Migrator::new(migrations_path).await.unwrap(); + migrator.run(&db.pool).await.unwrap(); + db + }); + + Self { + db: Some(db), + name, + url, + } + } + } + + pub fn db(&self) -> &Db { + self.db.as_ref().unwrap() + } + + fn truncate(&self) { + block_on(async { + let query = " + SELECT tablename FROM pg_tables + WHERE schemaname = 'public'; + "; + let table_names = sqlx::query_scalar::<_, String>(query) + .fetch_all(&self.db().pool) + .await + .unwrap(); + sqlx::query(&format!( + "TRUNCATE TABLE {} RESTART IDENTITY", + table_names.join(", ") + )) + .execute(&self.db().pool) + .await + .unwrap(); + }) + } + + fn teardown(mut self) { + let db = self.db.take().unwrap(); block_on(async { let query = " SELECT pg_terminate_backend(pg_stat_activity.pid) @@ -582,15 +616,27 @@ pub mod tests { "; sqlx::query(query) .bind(&self.name) - .execute(&self.db.pool) + .execute(&db.pool) .await .unwrap(); - self.db.pool.close().await; + db.pool.close().await; Postgres::drop_database(&self.url).await.unwrap(); }); } } + impl Drop for TestDb { + fn drop(&mut self) { + if let Some(db) = self.db.take() { + POOL.lock().push(TestDb { + db: Some(db), + name: mem::take(&mut self.name), + url: mem::take(&mut self.url), + }); + } + } + } + #[gpui::test] async fn test_get_users_by_ids() { let test_db = TestDb::new(); diff --git a/crates/server/src/rpc.rs b/crates/server/src/rpc.rs index 85b29540f5..bcf222f853 100644 --- a/crates/server/src/rpc.rs +++ b/crates/server/src/rpc.rs @@ -1196,8 +1196,9 @@ mod tests { #[cfg(test)] #[ctor::ctor] fn init_logger() { - // std::env::set_var("RUST_LOG", "info"); - env_logger::init(); + if std::env::var("RUST_LOG").is_ok() { + env_logger::init(); + } } #[gpui::test] diff --git a/crates/text/src/tests.rs b/crates/text/src/tests.rs index b4ac7d46d4..de735051ae 100644 --- a/crates/text/src/tests.rs +++ b/crates/text/src/tests.rs @@ -12,8 +12,9 @@ use util::test::Network; #[cfg(test)] #[ctor::ctor] fn init_logger() { - // std::env::set_var("RUST_LOG", "info"); - env_logger::init(); + if std::env::var("RUST_LOG").is_ok() { + env_logger::init(); + } } #[test] diff --git a/crates/zed/src/test.rs b/crates/zed/src/test.rs index a365fdc6f4..5bf43b555f 100644 --- a/crates/zed/src/test.rs +++ b/crates/zed/src/test.rs @@ -12,7 +12,9 @@ use workspace::Settings; #[cfg(test)] #[ctor::ctor] fn init_logger() { - env_logger::init(); + if std::env::var("RUST_LOG").is_ok() { + env_logger::init(); + } } pub fn test_app_state(cx: &mut MutableAppContext) -> Arc {