Cache DB connection in tests

To avoid creating a new DB connection on each call to the `db_client` method

Author: Kobzol

src/db/review_prefs.rs

@@ -117,12 +117,9 @@ mod tests {
     #[tokio::test]
     async fn insert_prefs_create_user() {
         run_db_test(|ctx| async {
-            let db = ctx.db_client().await;
-
             let user = user("Martin", 1);
-            upsert_review_prefs(&db, user.clone(), Some(1)).await?;
-
-            assert_eq!(get_user(&db, user.id).await?.unwrap(), user);
+            upsert_review_prefs(&ctx.db_client(), user.clone(), Some(1)).await?;
+            assert_eq!(get_user(&ctx.db_client(), user.id).await?.unwrap(), user);
 
             Ok(ctx)
         })
@@ -132,11 +129,12 @@ mod tests {
     #[tokio::test]
     async fn insert_max_assigned_prs() {
         run_db_test(|ctx| async {
-            let db = ctx.db_client().await;
-
-            upsert_review_prefs(&db, user("Martin", 1), Some(5)).await?;
+            upsert_review_prefs(&ctx.db_client(), user("Martin", 1), Some(5)).await?;
             assert_eq!(
-                get_review_prefs(&db, 1).await?.unwrap().max_assigned_prs,
+                get_review_prefs(&ctx.db_client(), 1)
+                    .await?
+                    .unwrap()
+                    .max_assigned_prs,
                 Some(5)
             );
 
@@ -148,7 +146,7 @@ mod tests {
     #[tokio::test]
     async fn update_max_assigned_prs() {
         run_db_test(|ctx| async {
-            let db = ctx.db_client().await;
+            let db = ctx.db_client();
 
             upsert_review_prefs(&db, user("Martin", 1), Some(5)).await?;
             assert_eq!(

src/db/users.rs

@@ -46,7 +46,7 @@ mod tests {
     #[tokio::test]
     async fn update_username_on_conflict() {
         run_db_test(|ctx| async {
-            let db = ctx.db_client().await;
+            let db = ctx.db_client();
 
             record_username(&db, 1, "Foo").await?;
             record_username(&db, 1, "Bar").await?;

src/handlers/pr_tracking.rs

@@ -43,6 +43,11 @@ impl ReviewerWorkqueue {
             .map(|prs| prs.len() as u64)
             .unwrap_or(0)
     }
+
+    #[cfg(test)]
+    pub fn set_user_prs(&mut self, user_id: UserId, prs: HashSet<PullRequestNumber>) {
+        self.reviewers.insert(user_id, prs);
+    }
 }
 
 pub(super) enum ReviewPrefsInput {

src/tests/mod.rs

@@ -8,18 +8,19 @@ use std::future::Future;
 use std::sync::Arc;
 use tokio::sync::RwLock;
 use tokio_postgres::config::Host;
-use tokio_postgres::{Config, GenericClient};
+use tokio_postgres::Config;
 
 pub(crate) mod github;
 
 /// Represents a connection to a Postgres database that can be
 /// used in integration tests to test logic that interacts with
 /// a database.
 pub(crate) struct TestContext {
-    pool: ClientPool,
     ctx: Context,
     db_name: String,
     original_db_url: String,
+    // Pre-cached client to avoid creating unnecessary connections in tests
+    client: PooledClient,
 }
 
 impl TestContext {
@@ -53,7 +54,8 @@ impl TestContext {
             db_name
         );
         let pool = ClientPool::new(test_db_url.clone());
-        db::run_migrations(&mut *pool.get().await)
+        let mut client = pool.get().await;
+        db::run_migrations(&mut client)
             .await
             .expect("Cannot run database migrations");
 
@@ -66,17 +68,17 @@ impl TestContext {
         );
         let ctx = Context {
             github,
-            db: ClientPool::new(test_db_url),
+            db: pool,
             username: "triagebot-test".to_string(),
             octocrab,
             workqueue: Arc::new(RwLock::new(Default::default())),
         };
 
         Self {
-            pool,
             db_name,
             original_db_url: db_url.to_string(),
             ctx,
+            client,
         }
     }
 
@@ -87,21 +89,21 @@ impl TestContext {
         &self.ctx
     }
 
-    pub(crate) async fn db_client(&self) -> PooledClient {
-        self.pool.get().await
+    pub(crate) fn db_client(&self) -> &PooledClient {
+        &self.client
     }
 
-    #[allow(dead_code)]
     pub(crate) async fn add_user(&self, name: &str, id: u64) {
-        record_username(self.db_client().await.client(), id, name)
+        record_username(self.db_client(), id, name)
             .await
             .expect("Cannot create user");
     }
 
     async fn finish(self) {
         // Cleanup the test database
         // First, we need to stop using the database
-        drop(self.pool);
+        drop(self.client);
+        drop(self.ctx);
 
         // Then we need to connect to the default database and drop our test DB
         let client = make_client(&self.original_db_url)