sentry - db - fix postgres test DB pool & test

Author: elpiel

sentry/src/db.rs

@@ -141,11 +141,8 @@ pub async fn setup_migrations(environment: &str) {
 #[cfg(test)]
 pub mod tests_postgres {
     use std::{
-        ops::Deref,
-        sync::{
-            atomic::{AtomicUsize, Ordering},
-            Arc,
-        },
+        ops::{Deref, DerefMut},
+        sync::atomic::{AtomicUsize, Ordering},
     };
 
     use deadpool::managed::{Manager as ManagerTrait, RecycleResult};
@@ -159,45 +156,64 @@ pub mod tests_postgres {
 
     pub type Pool = deadpool::managed::Pool<Manager>;
 
-    pub static DATABASE_POOL: Lazy<Pool> = Lazy::new(|| {
+    pub static DATABASE_POOL: Lazy<Pool> = Lazy::new(|| create_pool("test"));
+
+    /// we must have a duplication of the migration because of how migrant is handling migrations
+    /// we need to separately setup test migrations
+    pub static MIGRATIONS: &[&str] = &["20190806011140_initial-tables"];
+
+    fn create_pool(db_prefix: &str) -> Pool {
         let manager_config = ManagerConfig {
             recycling_method: deadpool_postgres::RecyclingMethod::Fast,
         };
-        let manager = Manager::new(POSTGRES_CONFIG.clone(), manager_config);
+        let manager = Manager::new(POSTGRES_CONFIG.clone(), manager_config, db_prefix);
 
         Pool::new(manager, 15)
-    });
-
-    /// we must have a duplication of the migration because of how migrant is handling migrations
-    /// we need to separately setup test migrations
-    pub static MIGRATIONS: &[&str] = &["20190806011140_initial-tables"];
+    }
 
     /// A Database is used to isolate test runs from each other
     /// we need to know the name of the database we've created.
     /// This will allow us the drop the database when we are recycling the connection
     pub struct Database {
-        inner: Arc<DatabaseInner>,
+        /// The database name that will be created by the pool `CREATE DATABASE`
+        /// This database will be set on configuration level of the underlying connection Pool for tests
+        pub name: String,
+        pub pool: deadpool_postgres::Pool<NoTls>,
     }
+
     impl Database {
         pub fn new(name: String, pool: DbPool) -> Self {
-            Self {
-                inner: Arc::new(DatabaseInner { name, pool }),
-            }
+            Self { name, pool }
         }
-    }
 
-    struct DatabaseInner {
-        /// The database name that will be created by the pool `CREATE DATABASE`
-        /// This database will be set on configuration level of the underlying connection Pool for tests
-        pub name: String,
-        pub pool: deadpool_postgres::Pool<NoTls>,
+        pub fn get_pool(&self) -> deadpool_postgres::Pool<NoTls> {
+            self.pool.clone()
+        }
     }
 
     impl Deref for Database {
         type Target = deadpool_postgres::Pool<NoTls>;
 
         fn deref(&self) -> &deadpool_postgres::Pool<NoTls> {
-            &self.inner.pool
+            &self.pool
+        }
+    }
+
+    impl DerefMut for Database {
+        fn deref_mut(&mut self) -> &mut Self::Target {
+            &mut self.pool
+        }
+    }
+
+    impl AsRef<deadpool_postgres::Pool<NoTls>> for Database {
+        fn as_ref(&self) -> &deadpool_postgres::Pool<NoTls> {
+            &self.pool
+        }
+    }
+
+    impl AsMut<deadpool_postgres::Pool<NoTls>> for Database {
+        fn as_mut(&mut self) -> &mut deadpool_postgres::Pool<NoTls> {
+            &mut self.pool
         }
     }
 
@@ -208,10 +224,15 @@ pub mod tests_postgres {
         base_pool: deadpool_postgres::Pool<NoTls>,
         manager_config: ManagerConfig,
         index: AtomicUsize,
+        db_prefix: String,
     }
 
     impl Manager {
-        pub fn new(base_config: tokio_postgres::Config, manager_config: ManagerConfig) -> Self {
+        pub fn new(
+            base_config: tokio_postgres::Config,
+            manager_config: ManagerConfig,
+            db_prefix: &str,
+        ) -> Self {
             // We need to create the schema with a temporary connection, in order to use it for the real Test Pool
             let base_manager = deadpool_postgres::Manager::from_config(
                 base_config.clone(),
@@ -220,19 +241,21 @@ pub mod tests_postgres {
             );
             let base_pool = deadpool_postgres::Pool::new(base_manager, 15);
 
-            Self::new_with_pool(base_pool, base_config, manager_config)
+            Self::new_with_pool(base_pool, base_config, manager_config, db_prefix)
         }
 
         pub fn new_with_pool(
             base_pool: deadpool_postgres::Pool<NoTls>,
             base_config: tokio_postgres::Config,
             manager_config: ManagerConfig,
+            db_prefix: &str,
         ) -> Self {
             Self {
                 base_config,
                 base_pool,
                 manager_config,
                 index: AtomicUsize::new(0),
+                db_prefix: db_prefix.into(),
             }
         }
     }
@@ -245,7 +268,9 @@ pub mod tests_postgres {
 
         async fn create(&self) -> Result<Self::Type, Self::Error> {
             let pool_index = self.index.fetch_add(1, Ordering::SeqCst);
-            let db_name = format!("test_{}", pool_index);
+
+            // e.g. test_0, test_1, test_2
+            let db_name = format!("{}_{}", self.db_prefix, pool_index);
 
             // 1. Drop the database if it exists - if a test failed before, the database wouldn't have been removed
             // 2. Create database
@@ -279,16 +304,18 @@ pub mod tests_postgres {
         }
 
         async fn recycle(&self, database: &mut Database) -> RecycleResult<Self::Error> {
-            let queries = format!("DROP DATABASE {0} WITH (FORCE);", database.inner.name);
-            let result = self
-                .base_pool
+            // DROP the public schema and create it again for usage after recycling
+            let queries = format!("DROP SCHEMA public CASCADE; CREATE SCHEMA public;");
+            let result = database
+                .pool
                 .get()
                 .await?
                 .simple_query(&queries)
                 .await
                 .map_err(|err| PoolError::Backend(err))?;
-            assert_eq!(1, result.len());
+            assert_eq!(2, result.len());
             assert!(matches!(result[0], SimpleQueryMessage::CommandComplete(..)));
+            assert!(matches!(result[1], SimpleQueryMessage::CommandComplete(..)));
 
             Ok(())
         }
@@ -318,6 +345,59 @@ pub mod tests_postgres {
 
         Ok(client.batch_execute(&full_query).await?)
     }
+
+    #[cfg(test)]
+    mod test {
+        use super::*;
+
+        #[tokio::test]
+        /// Does not use the `DATABASE_POOL` as other tests can interfere with the pool objects!
+        async fn test_postgres_pool() {
+            let pool = create_pool("testing_pool");
+
+            let database_1 = pool.get().await.expect("Should get");
+            let status = pool.status();
+            assert_eq!(status.size, 1);
+            assert_eq!(status.available, 0);
+
+            let database_2 = pool.get().await.expect("Should get");
+            let status = pool.status();
+            assert_eq!(status.size, 2);
+            assert_eq!(status.available, 0);
+
+            drop(database_1);
+            let status = pool.status();
+            assert_eq!(status.size, 2);
+            assert_eq!(status.available, 1);
+
+            drop(database_2);
+            let status = pool.status();
+            assert_eq!(status.size, 2);
+            assert_eq!(status.available, 2);
+
+            let database_3 = pool.get().await.expect("Should get");
+            let status = pool.status();
+            assert_eq!(status.size, 2);
+            assert_eq!(status.available, 1);
+
+            let database_4 = pool.get().await.expect("Should get");
+            let status = pool.status();
+            assert_eq!(status.size, 2);
+            assert_eq!(status.available, 0);
+
+            let database_5 = pool.get().await.expect("Should get");
+            let status = pool.status();
+            assert_eq!(status.size, 3);
+            assert_eq!(status.available, 0);
+
+            drop(database_3);
+            drop(database_4);
+            drop(database_5);
+            let status = pool.status();
+            assert_eq!(status.size, 3);
+            assert_eq!(status.available, 3);
+        }
+    }
 }
 
 #[cfg(test)]

sentry/src/db/campaign.rs

@@ -51,28 +51,28 @@ pub async fn fetch_campaign(pool: DbPool, campaign: &Campaign) -> Result<Campaig
 
 #[cfg(test)]
 mod test {
-    use primitives::{campaign::Campaign, util::tests::prep_db::DUMMY_CAMPAIGN};
+    use primitives::util::tests::prep_db::DUMMY_CAMPAIGN;
 
     use crate::db::tests_postgres::{setup_test_migrations, DATABASE_POOL};
 
     use super::*;
 
     #[tokio::test]
     async fn it_inserts_and_fetches_campaign() {
-        let db_pool = DATABASE_POOL.get().await.expect("Should get a DB pool");
+        let database = DATABASE_POOL.get().await.expect("Should get a DB pool");
 
-        setup_test_migrations(db_pool.clone())
+        setup_test_migrations(database.get_pool())
             .await
             .expect("Migrations should succeed");
 
         let campaign_for_testing = DUMMY_CAMPAIGN.clone();
-        let is_inserted = insert_campaign(&db_pool.clone(), &campaign_for_testing)
+        let is_inserted = insert_campaign(&database.get_pool(), &campaign_for_testing)
             .await
             .expect("Should succeed");
 
         assert!(is_inserted);
 
-        let fetched_campaign: Campaign = fetch_campaign(db_pool.clone(), &campaign_for_testing)
+        let fetched_campaign = fetch_campaign(database.get_pool(), &campaign_for_testing)
             .await
             .expect("Should fetch successfully");
 

sentry/src/db/spendable.rs

@@ -60,9 +60,9 @@ mod test {
 
     #[tokio::test]
     async fn it_inserts_and_fetches_spendable() {
-        let db_pool = DATABASE_POOL.get().await.expect("Should get a DB pool");
+        let database = DATABASE_POOL.get().await.expect("Should get a DB pool");
 
-        setup_test_migrations(db_pool.clone())
+        setup_test_migrations(database.get_pool())
             .await
             .expect("Migrations should succeed");
 
@@ -74,16 +74,19 @@ mod test {
                 still_on_create2: UnifiedNum::from(500_000),
             },
         };
-        let is_inserted = insert_spendable(db_pool.clone(), &spendable)
+        let is_inserted = insert_spendable(database.get_pool(), &spendable)
             .await
             .expect("Should succeed");
 
         assert!(is_inserted);
 
-        let fetched_spendable =
-            fetch_spendable(db_pool.clone(), &spendable.spender, &spendable.channel.id())
-                .await
-                .expect("Should fetch successfully");
+        let fetched_spendable = fetch_spendable(
+            database.get_pool(),
+            &spendable.spender,
+            &spendable.channel.id(),
+        )
+        .await
+        .expect("Should fetch successfully");
 
         assert_eq!(spendable, fetched_spendable);
     }