Use a real database in tests, but block on db calls

Max Brunsfeld and Nathan Sobo created

Co-authored-by: Nathan Sobo <nathan@zed.dev>

Change summary

crates/collab/src/db.rs                | 926 ++++++++++++++-------------
crates/collab/src/integration_tests.rs |  17 
2 files changed, 508 insertions(+), 435 deletions(-)

Detailed changes

crates/collab/src/db.rs 🔗

@@ -188,6 +188,20 @@ pub struct PostgresDb {
     pool: sqlx::PgPool,
 }
 
+macro_rules! test_support {
+    ($self:ident, { $($token:tt)* }) => {{
+        let body = async {
+            $($token)*
+        };
+
+        if cfg!(test) {
+            tokio::runtime::Builder::new_current_thread().enable_io().enable_time().build().unwrap().block_on(body)
+        } else {
+            body.await
+        }
+    }};
+}
+
 impl PostgresDb {
     pub async fn new(url: &str, max_connections: u32) -> Result<Self> {
         let pool = DbOptions::new()
@@ -262,51 +276,58 @@ impl Db for PostgresDb {
         admin: bool,
         params: NewUserParams,
     ) -> Result<NewUserResult> {
-        let query = "
-            INSERT INTO users (email_address, github_login, github_user_id, admin)
-            VALUES ($1, $2, $3, $4)
-            ON CONFLICT (github_login) DO UPDATE SET github_login = excluded.github_login
-            RETURNING id, metrics_id::text
-        ";
-        let (user_id, metrics_id): (UserId, String) = sqlx::query_as(query)
-            .bind(email_address)
-            .bind(params.github_login)
-            .bind(params.github_user_id)
-            .bind(admin)
-            .fetch_one(&self.pool)
-            .await?;
-        Ok(NewUserResult {
-            user_id,
-            metrics_id,
-            signup_device_id: None,
-            inviting_user_id: None,
+        test_support!(self, {
+            let query = "
+                INSERT INTO users (email_address, github_login, github_user_id, admin)
+                VALUES ($1, $2, $3, $4)
+                ON CONFLICT (github_login) DO UPDATE SET github_login = excluded.github_login
+                RETURNING id, metrics_id::text
+            ";
+
+            let (user_id, metrics_id): (UserId, String) = sqlx::query_as(query)
+                .bind(email_address)
+                .bind(params.github_login)
+                .bind(params.github_user_id)
+                .bind(admin)
+                .fetch_one(&self.pool)
+                .await?;
+            Ok(NewUserResult {
+                user_id,
+                metrics_id,
+                signup_device_id: None,
+                inviting_user_id: None,
+            })
         })
     }
 
     async fn get_all_users(&self, page: u32, limit: u32) -> Result<Vec<User>> {
-        let query = "SELECT * FROM users ORDER BY github_login ASC LIMIT $1 OFFSET $2";
-        Ok(sqlx::query_as(query)
-            .bind(limit as i32)
-            .bind((page * limit) as i32)
-            .fetch_all(&self.pool)
-            .await?)
+        test_support!(self, {
+            let query = "SELECT * FROM users ORDER BY github_login ASC LIMIT $1 OFFSET $2";
+            Ok(sqlx::query_as(query)
+                .bind(limit as i32)
+                .bind((page * limit) as i32)
+                .fetch_all(&self.pool)
+                .await?)
+        })
     }
 
     async fn fuzzy_search_users(&self, name_query: &str, limit: u32) -> Result<Vec<User>> {
-        let like_string = Self::fuzzy_like_string(name_query);
-        let query = "
-            SELECT users.*
-            FROM users
-            WHERE github_login ILIKE $1
-            ORDER BY github_login <-> $2
-            LIMIT $3
-        ";
-        Ok(sqlx::query_as(query)
-            .bind(like_string)
-            .bind(name_query)
-            .bind(limit as i32)
-            .fetch_all(&self.pool)
-            .await?)
+        test_support!(self, {
+            let like_string = Self::fuzzy_like_string(name_query);
+            let query = "
+                SELECT users.*
+                FROM users
+                WHERE github_login ILIKE $1
+                ORDER BY github_login <-> $2
+                LIMIT $3
+            ";
+            Ok(sqlx::query_as(query)
+                .bind(like_string)
+                .bind(name_query)
+                .bind(limit as i32)
+                .fetch_all(&self.pool)
+                .await?)
+        })
     }
 
     async fn get_user_by_id(&self, id: UserId) -> Result<Option<User>> {
@@ -315,42 +336,48 @@ impl Db for PostgresDb {
     }
 
     async fn get_user_metrics_id(&self, id: UserId) -> Result<String> {
-        let query = "
-            SELECT metrics_id::text
-            FROM users
-            WHERE id = $1
-        ";
-        Ok(sqlx::query_scalar(query)
-            .bind(id)
-            .fetch_one(&self.pool)
-            .await?)
+        test_support!(self, {
+            let query = "
+                SELECT metrics_id::text
+                FROM users
+                WHERE id = $1
+            ";
+            Ok(sqlx::query_scalar(query)
+                .bind(id)
+                .fetch_one(&self.pool)
+                .await?)
+        })
     }
 
     async fn get_users_by_ids(&self, ids: Vec<UserId>) -> Result<Vec<User>> {
-        let ids = ids.into_iter().map(|id| id.0).collect::<Vec<_>>();
-        let query = "
-            SELECT users.*
-            FROM users
-            WHERE users.id = ANY ($1)
-        ";
-        Ok(sqlx::query_as(query)
-            .bind(&ids)
-            .fetch_all(&self.pool)
-            .await?)
+        test_support!(self, {
+            let ids = ids.into_iter().map(|id| id.0).collect::<Vec<_>>();
+            let query = "
+                SELECT users.*
+                FROM users
+                WHERE users.id = ANY ($1)
+            ";
+            Ok(sqlx::query_as(query)
+                .bind(&ids)
+                .fetch_all(&self.pool)
+                .await?)
+        })
     }
 
     async fn get_users_with_no_invites(&self, invited_by_another_user: bool) -> Result<Vec<User>> {
-        let query = format!(
-            "
-            SELECT users.*
-            FROM users
-            WHERE invite_count = 0
-            AND inviter_id IS{} NULL
-            ",
-            if invited_by_another_user { " NOT" } else { "" }
-        );
-
-        Ok(sqlx::query_as(&query).fetch_all(&self.pool).await?)
+        test_support!(self, {
+            let query = format!(
+                "
+                SELECT users.*
+                FROM users
+                WHERE invite_count = 0
+                AND inviter_id IS{} NULL
+                ",
+                if invited_by_another_user { " NOT" } else { "" }
+            );
+
+            Ok(sqlx::query_as(&query).fetch_all(&self.pool).await?)
+        })
     }
 
     async fn get_user_by_github_account(
@@ -358,176 +385,193 @@ impl Db for PostgresDb {
         github_login: &str,
         github_user_id: Option<i32>,
     ) -> Result<Option<User>> {
-        if let Some(github_user_id) = github_user_id {
-            let mut user = sqlx::query_as::<_, User>(
-                "
-                UPDATE users
-                SET github_login = $1
-                WHERE github_user_id = $2
-                RETURNING *
-                ",
-            )
-            .bind(github_login)
-            .bind(github_user_id)
-            .fetch_optional(&self.pool)
-            .await?;
-
-            if user.is_none() {
-                user = sqlx::query_as::<_, User>(
+        test_support!(self, {
+            if let Some(github_user_id) = github_user_id {
+                let mut user = sqlx::query_as::<_, User>(
                     "
                     UPDATE users
-                    SET github_user_id = $1
-                    WHERE github_login = $2
+                    SET github_login = $1
+                    WHERE github_user_id = $2
                     RETURNING *
                     ",
                 )
+                .bind(github_login)
                 .bind(github_user_id)
+                .fetch_optional(&self.pool)
+                .await?;
+
+                if user.is_none() {
+                    user = sqlx::query_as::<_, User>(
+                        "
+                        UPDATE users
+                        SET github_user_id = $1
+                        WHERE github_login = $2
+                        RETURNING *
+                        ",
+                    )
+                    .bind(github_user_id)
+                    .bind(github_login)
+                    .fetch_optional(&self.pool)
+                    .await?;
+                }
+
+                Ok(user)
+            } else {
+                let user = sqlx::query_as(
+                    "
+                    SELECT * FROM users
+                    WHERE github_login = $1
+                    LIMIT 1
+                    ",
+                )
                 .bind(github_login)
                 .fetch_optional(&self.pool)
                 .await?;
+                Ok(user)
             }
-
-            Ok(user)
-        } else {
-            Ok(sqlx::query_as(
-                "
-                SELECT * FROM users
-                WHERE github_login = $1
-                LIMIT 1
-                ",
-            )
-            .bind(github_login)
-            .fetch_optional(&self.pool)
-            .await?)
-        }
+        })
     }
 
     async fn set_user_is_admin(&self, id: UserId, is_admin: bool) -> Result<()> {
-        let query = "UPDATE users SET admin = $1 WHERE id = $2";
-        Ok(sqlx::query(query)
-            .bind(is_admin)
-            .bind(id.0)
-            .execute(&self.pool)
-            .await
-            .map(drop)?)
+        test_support!(self, {
+            let query = "UPDATE users SET admin = $1 WHERE id = $2";
+            Ok(sqlx::query(query)
+                .bind(is_admin)
+                .bind(id.0)
+                .execute(&self.pool)
+                .await
+                .map(drop)?)
+        })
     }
 
     async fn set_user_connected_once(&self, id: UserId, connected_once: bool) -> Result<()> {
-        let query = "UPDATE users SET connected_once = $1 WHERE id = $2";
-        Ok(sqlx::query(query)
-            .bind(connected_once)
-            .bind(id.0)
-            .execute(&self.pool)
-            .await
-            .map(drop)?)
+        test_support!(self, {
+            let query = "UPDATE users SET connected_once = $1 WHERE id = $2";
+            Ok(sqlx::query(query)
+                .bind(connected_once)
+                .bind(id.0)
+                .execute(&self.pool)
+                .await
+                .map(drop)?)
+        })
     }
 
     async fn destroy_user(&self, id: UserId) -> Result<()> {
-        let query = "DELETE FROM access_tokens WHERE user_id = $1;";
-        sqlx::query(query)
-            .bind(id.0)
-            .execute(&self.pool)
-            .await
-            .map(drop)?;
-        let query = "DELETE FROM users WHERE id = $1;";
-        Ok(sqlx::query(query)
-            .bind(id.0)
-            .execute(&self.pool)
-            .await
-            .map(drop)?)
+        test_support!(self, {
+            let query = "DELETE FROM access_tokens WHERE user_id = $1;";
+            sqlx::query(query)
+                .bind(id.0)
+                .execute(&self.pool)
+                .await
+                .map(drop)?;
+            let query = "DELETE FROM users WHERE id = $1;";
+            Ok(sqlx::query(query)
+                .bind(id.0)
+                .execute(&self.pool)
+                .await
+                .map(drop)?)
+        })
     }
 
     // signups
 
     async fn create_signup(&self, signup: Signup) -> Result<()> {
-        sqlx::query(
-            "
-            INSERT INTO signups
-            (
-                email_address,
-                email_confirmation_code,
-                email_confirmation_sent,
-                platform_linux,
-                platform_mac,
-                platform_windows,
-                platform_unknown,
-                editor_features,
-                programming_languages,
-                device_id
+        test_support!(self, {
+            sqlx::query(
+                "
+                INSERT INTO signups
+                (
+                    email_address,
+                    email_confirmation_code,
+                    email_confirmation_sent,
+                    platform_linux,
+                    platform_mac,
+                    platform_windows,
+                    platform_unknown,
+                    editor_features,
+                    programming_languages,
+                    device_id
+                )
+                VALUES
+                    ($1, $2, 'f', $3, $4, $5, 'f', $6, $7, $8)
+                RETURNING id
+                ",
             )
-            VALUES
-                ($1, $2, 'f', $3, $4, $5, 'f', $6, $7, $8)
-            RETURNING id
-            ",
-        )
-        .bind(&signup.email_address)
-        .bind(&random_email_confirmation_code())
-        .bind(&signup.platform_linux)
-        .bind(&signup.platform_mac)
-        .bind(&signup.platform_windows)
-        .bind(&signup.editor_features)
-        .bind(&signup.programming_languages)
-        .bind(&signup.device_id)
-        .execute(&self.pool)
-        .await?;
-        Ok(())
+            .bind(&signup.email_address)
+            .bind(&random_email_confirmation_code())
+            .bind(&signup.platform_linux)
+            .bind(&signup.platform_mac)
+            .bind(&signup.platform_windows)
+            .bind(&signup.editor_features)
+            .bind(&signup.programming_languages)
+            .bind(&signup.device_id)
+            .execute(&self.pool)
+            .await?;
+            Ok(())
+        })
     }
 
     async fn get_waitlist_summary(&self) -> Result<WaitlistSummary> {
-        Ok(sqlx::query_as(
-            "
-            SELECT
-                COUNT(*) as count,
-                COALESCE(SUM(CASE WHEN platform_linux THEN 1 ELSE 0 END), 0) as linux_count,
-                COALESCE(SUM(CASE WHEN platform_mac THEN 1 ELSE 0 END), 0) as mac_count,
-                COALESCE(SUM(CASE WHEN platform_windows THEN 1 ELSE 0 END), 0) as windows_count,
-                COALESCE(SUM(CASE WHEN platform_unknown THEN 1 ELSE 0 END), 0) as unknown_count
-            FROM (
-                SELECT *
-                FROM signups
-                WHERE
-                    NOT email_confirmation_sent
-            ) AS unsent
-            ",
-        )
-        .fetch_one(&self.pool)
-        .await?)
+        test_support!(self, {
+            Ok(sqlx::query_as(
+                "
+                SELECT
+                    COUNT(*) as count,
+                    COALESCE(SUM(CASE WHEN platform_linux THEN 1 ELSE 0 END), 0) as linux_count,
+                    COALESCE(SUM(CASE WHEN platform_mac THEN 1 ELSE 0 END), 0) as mac_count,
+                    COALESCE(SUM(CASE WHEN platform_windows THEN 1 ELSE 0 END), 0) as windows_count,
+                    COALESCE(SUM(CASE WHEN platform_unknown THEN 1 ELSE 0 END), 0) as unknown_count
+                FROM (
+                    SELECT *
+                    FROM signups
+                    WHERE
+                        NOT email_confirmation_sent
+                ) AS unsent
+                ",
+            )
+            .fetch_one(&self.pool)
+            .await?)
+        })
     }
 
     async fn get_unsent_invites(&self, count: usize) -> Result<Vec<Invite>> {
-        Ok(sqlx::query_as(
-            "
-            SELECT
-                email_address, email_confirmation_code
-            FROM signups
-            WHERE
-                NOT email_confirmation_sent AND
-                (platform_mac OR platform_unknown)
-            LIMIT $1
-            ",
-        )
-        .bind(count as i32)
-        .fetch_all(&self.pool)
-        .await?)
+        test_support!(self, {
+            Ok(sqlx::query_as(
+                "
+                SELECT
+                    email_address, email_confirmation_code
+                FROM signups
+                WHERE
+                    NOT email_confirmation_sent AND
+                    (platform_mac OR platform_unknown)
+                LIMIT $1
+                ",
+            )
+            .bind(count as i32)
+            .fetch_all(&self.pool)
+            .await?)
+        })
     }
 
     async fn record_sent_invites(&self, invites: &[Invite]) -> Result<()> {
-        sqlx::query(
-            "
-            UPDATE signups
-            SET email_confirmation_sent = 't'
-            WHERE email_address = ANY ($1)
-            ",
-        )
-        .bind(
-            &invites
-                .iter()
-                .map(|s| s.email_address.as_str())
-                .collect::<Vec<_>>(),
-        )
-        .execute(&self.pool)
-        .await?;
-        Ok(())
+        test_support!(self, {
+            sqlx::query(
+                "
+                UPDATE signups
+                SET email_confirmation_sent = 't'
+                WHERE email_address = ANY ($1)
+                ",
+            )
+            .bind(
+                &invites
+                    .iter()
+                    .map(|s| s.email_address.as_str())
+                    .collect::<Vec<_>>(),
+            )
+            .execute(&self.pool)
+            .await?;
+            Ok(())
+        })
     }
 
     async fn create_user_from_invite(
@@ -535,176 +579,184 @@ impl Db for PostgresDb {
         invite: &Invite,
         user: NewUserParams,
     ) -> Result<Option<NewUserResult>> {
-        let mut tx = self.pool.begin().await?;
-
-        let (signup_id, existing_user_id, inviting_user_id, signup_device_id): (
-            i32,
-            Option<UserId>,
-            Option<UserId>,
-            Option<String>,
-        ) = sqlx::query_as(
-            "
-            SELECT id, user_id, inviting_user_id, device_id
-            FROM signups
-            WHERE
-                email_address = $1 AND
-                email_confirmation_code = $2
-            ",
-        )
-        .bind(&invite.email_address)
-        .bind(&invite.email_confirmation_code)
-        .fetch_optional(&mut tx)
-        .await?
-        .ok_or_else(|| Error::Http(StatusCode::NOT_FOUND, "no such invite".to_string()))?;
-
-        if existing_user_id.is_some() {
-            return Ok(None);
-        }
+        test_support!(self, {
+            let mut tx = self.pool.begin().await?;
 
-        let (user_id, metrics_id): (UserId, String) = sqlx::query_as(
-            "
-            INSERT INTO users
-            (email_address, github_login, github_user_id, admin, invite_count, invite_code)
-            VALUES
-            ($1, $2, $3, 'f', $4, $5)
-            ON CONFLICT (github_login) DO UPDATE SET
-                email_address = excluded.email_address,
-                github_user_id = excluded.github_user_id,
-                admin = excluded.admin
-            RETURNING id, metrics_id::text
-            ",
-        )
-        .bind(&invite.email_address)
-        .bind(&user.github_login)
-        .bind(&user.github_user_id)
-        .bind(&user.invite_count)
-        .bind(random_invite_code())
-        .fetch_one(&mut tx)
-        .await?;
-
-        sqlx::query(
-            "
-            UPDATE signups
-            SET user_id = $1
-            WHERE id = $2
-            ",
-        )
-        .bind(&user_id)
-        .bind(&signup_id)
-        .execute(&mut tx)
-        .await?;
-
-        if let Some(inviting_user_id) = inviting_user_id {
-            let id: Option<UserId> = sqlx::query_scalar(
+            let (signup_id, existing_user_id, inviting_user_id, signup_device_id): (
+                i32,
+                Option<UserId>,
+                Option<UserId>,
+                Option<String>,
+            ) = sqlx::query_as(
                 "
-                UPDATE users
-                SET invite_count = invite_count - 1
-                WHERE id = $1 AND invite_count > 0
-                RETURNING id
+                SELECT id, user_id, inviting_user_id, device_id
+                FROM signups
+                WHERE
+                    email_address = $1 AND
+                    email_confirmation_code = $2
                 ",
             )
-            .bind(&inviting_user_id)
+            .bind(&invite.email_address)
+            .bind(&invite.email_confirmation_code)
             .fetch_optional(&mut tx)
-            .await?;
+            .await?
+            .ok_or_else(|| Error::Http(StatusCode::NOT_FOUND, "no such invite".to_string()))?;
 
-            if id.is_none() {
-                Err(Error::Http(
-                    StatusCode::UNAUTHORIZED,
-                    "no invites remaining".to_string(),
-                ))?;
+            if existing_user_id.is_some() {
+                return Ok(None);
             }
 
-            sqlx::query(
+            let (user_id, metrics_id): (UserId, String) = sqlx::query_as(
                 "
-                INSERT INTO contacts
-                    (user_id_a, user_id_b, a_to_b, should_notify, accepted)
+                INSERT INTO users
+                (email_address, github_login, github_user_id, admin, invite_count, invite_code)
                 VALUES
-                    ($1, $2, 't', 't', 't')
-                ON CONFLICT DO NOTHING
+                ($1, $2, $3, 'f', $4, $5)
+                ON CONFLICT (github_login) DO UPDATE SET
+                    email_address = excluded.email_address,
+                    github_user_id = excluded.github_user_id,
+                    admin = excluded.admin
+                RETURNING id, metrics_id::text
                 ",
             )
-            .bind(inviting_user_id)
-            .bind(user_id)
+            .bind(&invite.email_address)
+            .bind(&user.github_login)
+            .bind(&user.github_user_id)
+            .bind(&user.invite_count)
+            .bind(random_invite_code())
+            .fetch_one(&mut tx)
+            .await?;
+
+            sqlx::query(
+                "
+                UPDATE signups
+                SET user_id = $1
+                WHERE id = $2
+                ",
+            )
+            .bind(&user_id)
+            .bind(&signup_id)
             .execute(&mut tx)
             .await?;
-        }
 
-        tx.commit().await?;
-        Ok(Some(NewUserResult {
-            user_id,
-            metrics_id,
-            inviting_user_id,
-            signup_device_id,
-        }))
+            if let Some(inviting_user_id) = inviting_user_id {
+                let id: Option<UserId> = sqlx::query_scalar(
+                    "
+                    UPDATE users
+                    SET invite_count = invite_count - 1
+                    WHERE id = $1 AND invite_count > 0
+                    RETURNING id
+                    ",
+                )
+                .bind(&inviting_user_id)
+                .fetch_optional(&mut tx)
+                .await?;
+
+                if id.is_none() {
+                    Err(Error::Http(
+                        StatusCode::UNAUTHORIZED,
+                        "no invites remaining".to_string(),
+                    ))?;
+                }
+
+                sqlx::query(
+                    "
+                    INSERT INTO contacts
+                        (user_id_a, user_id_b, a_to_b, should_notify, accepted)
+                    VALUES
+                        ($1, $2, 't', 't', 't')
+                    ON CONFLICT DO NOTHING
+                    ",
+                )
+                .bind(inviting_user_id)
+                .bind(user_id)
+                .execute(&mut tx)
+                .await?;
+            }
+
+            tx.commit().await?;
+            Ok(Some(NewUserResult {
+                user_id,
+                metrics_id,
+                inviting_user_id,
+                signup_device_id,
+            }))
+        })
     }
 
     // invite codes
 
     async fn set_invite_count_for_user(&self, id: UserId, count: u32) -> Result<()> {
-        let mut tx = self.pool.begin().await?;
-        if count > 0 {
+        test_support!(self, {
+            let mut tx = self.pool.begin().await?;
+            if count > 0 {
+                sqlx::query(
+                    "
+                    UPDATE users
+                    SET invite_code = $1
+                    WHERE id = $2 AND invite_code IS NULL
+                ",
+                )
+                .bind(random_invite_code())
+                .bind(id)
+                .execute(&mut tx)
+                .await?;
+            }
+
             sqlx::query(
                 "
                 UPDATE users
-                SET invite_code = $1
-                WHERE id = $2 AND invite_code IS NULL
-            ",
+                SET invite_count = $1
+                WHERE id = $2
+                ",
             )
-            .bind(random_invite_code())
+            .bind(count as i32)
             .bind(id)
             .execute(&mut tx)
             .await?;
-        }
-
-        sqlx::query(
-            "
-            UPDATE users
-            SET invite_count = $1
-            WHERE id = $2
-            ",
-        )
-        .bind(count as i32)
-        .bind(id)
-        .execute(&mut tx)
-        .await?;
-        tx.commit().await?;
-        Ok(())
+            tx.commit().await?;
+            Ok(())
+        })
     }
 
     async fn get_invite_code_for_user(&self, id: UserId) -> Result<Option<(String, u32)>> {
-        let result: Option<(String, i32)> = sqlx::query_as(
-            "
-                SELECT invite_code, invite_count
-                FROM users
-                WHERE id = $1 AND invite_code IS NOT NULL 
-            ",
-        )
-        .bind(id)
-        .fetch_optional(&self.pool)
-        .await?;
-        if let Some((code, count)) = result {
-            Ok(Some((code, count.try_into().map_err(anyhow::Error::new)?)))
-        } else {
-            Ok(None)
-        }
+        test_support!(self, {
+            let result: Option<(String, i32)> = sqlx::query_as(
+                "
+                    SELECT invite_code, invite_count
+                    FROM users
+                    WHERE id = $1 AND invite_code IS NOT NULL 
+                ",
+            )
+            .bind(id)
+            .fetch_optional(&self.pool)
+            .await?;
+            if let Some((code, count)) = result {
+                Ok(Some((code, count.try_into().map_err(anyhow::Error::new)?)))
+            } else {
+                Ok(None)
+            }
+        })
     }
 
     async fn get_user_for_invite_code(&self, code: &str) -> Result<User> {
-        sqlx::query_as(
-            "
-                SELECT *
-                FROM users
-                WHERE invite_code = $1
-            ",
-        )
-        .bind(code)
-        .fetch_optional(&self.pool)
-        .await?
-        .ok_or_else(|| {
-            Error::Http(
-                StatusCode::NOT_FOUND,
-                "that invite code does not exist".to_string(),
+        test_support!(self, {
+            sqlx::query_as(
+                "
+                    SELECT *
+                    FROM users
+                    WHERE invite_code = $1
+                ",
             )
+            .bind(code)
+            .fetch_optional(&self.pool)
+            .await?
+            .ok_or_else(|| {
+                Error::Http(
+                    StatusCode::NOT_FOUND,
+                    "that invite code does not exist".to_string(),
+                )
+            })
         })
     }
 
@@ -714,113 +766,119 @@ impl Db for PostgresDb {
         email_address: &str,
         device_id: Option<&str>,
     ) -> Result<Invite> {
-        let mut tx = self.pool.begin().await?;
-
-        let existing_user: Option<UserId> = sqlx::query_scalar(
-            "
-            SELECT id
-            FROM users
-            WHERE email_address = $1
-            ",
-        )
-        .bind(email_address)
-        .fetch_optional(&mut tx)
-        .await?;
-        if existing_user.is_some() {
-            Err(anyhow!("email address is already in use"))?;
-        }
+        test_support!(self, {
+            let mut tx = self.pool.begin().await?;
 
-        let row: Option<(UserId, i32)> = sqlx::query_as(
-            "
-            SELECT id, invite_count
-            FROM users
-            WHERE invite_code = $1
-            ",
-        )
-        .bind(code)
-        .fetch_optional(&mut tx)
-        .await?;
-
-        let (inviter_id, invite_count) = match row {
-            Some(row) => row,
-            None => Err(Error::Http(
-                StatusCode::NOT_FOUND,
-                "invite code not found".to_string(),
-            ))?,
-        };
+            let existing_user: Option<UserId> = sqlx::query_scalar(
+                "
+                SELECT id
+                FROM users
+                WHERE email_address = $1
+                ",
+            )
+            .bind(email_address)
+            .fetch_optional(&mut tx)
+            .await?;
+            if existing_user.is_some() {
+                Err(anyhow!("email address is already in use"))?;
+            }
 
-        if invite_count == 0 {
-            Err(Error::Http(
-                StatusCode::UNAUTHORIZED,
-                "no invites remaining".to_string(),
-            ))?;
-        }
+            let row: Option<(UserId, i32)> = sqlx::query_as(
+                "
+                SELECT id, invite_count
+                FROM users
+                WHERE invite_code = $1
+                ",
+            )
+            .bind(code)
+            .fetch_optional(&mut tx)
+            .await?;
 
-        let email_confirmation_code: String = sqlx::query_scalar(
-            "
-            INSERT INTO signups
-            (
-                email_address,
-                email_confirmation_code,
-                email_confirmation_sent,
-                inviting_user_id,
-                platform_linux,
-                platform_mac,
-                platform_windows,
-                platform_unknown,
-                device_id
+            let (inviter_id, invite_count) = match row {
+                Some(row) => row,
+                None => Err(Error::Http(
+                    StatusCode::NOT_FOUND,
+                    "invite code not found".to_string(),
+                ))?,
+            };
+
+            if invite_count == 0 {
+                Err(Error::Http(
+                    StatusCode::UNAUTHORIZED,
+                    "no invites remaining".to_string(),
+                ))?;
+            }
+
+            let email_confirmation_code: String = sqlx::query_scalar(
+                "
+                INSERT INTO signups
+                (
+                    email_address,
+                    email_confirmation_code,
+                    email_confirmation_sent,
+                    inviting_user_id,
+                    platform_linux,
+                    platform_mac,
+                    platform_windows,
+                    platform_unknown,
+                    device_id
+                )
+                VALUES
+                    ($1, $2, 'f', $3, 'f', 'f', 'f', 't', $4)
+                ON CONFLICT (email_address)
+                DO UPDATE SET
+                    inviting_user_id = excluded.inviting_user_id
+                RETURNING email_confirmation_code
+                ",
             )
-            VALUES
-                ($1, $2, 'f', $3, 'f', 'f', 'f', 't', $4)
-            ON CONFLICT (email_address)
-            DO UPDATE SET
-                inviting_user_id = excluded.inviting_user_id
-            RETURNING email_confirmation_code
-            ",
-        )
-        .bind(&email_address)
-        .bind(&random_email_confirmation_code())
-        .bind(&inviter_id)
-        .bind(&device_id)
-        .fetch_one(&mut tx)
-        .await?;
-
-        tx.commit().await?;
-
-        Ok(Invite {
-            email_address: email_address.into(),
-            email_confirmation_code,
+            .bind(&email_address)
+            .bind(&random_email_confirmation_code())
+            .bind(&inviter_id)
+            .bind(&device_id)
+            .fetch_one(&mut tx)
+            .await?;
+
+            tx.commit().await?;
+
+            Ok(Invite {
+                email_address: email_address.into(),
+                email_confirmation_code,
+            })
         })
     }
 
     // projects
 
     async fn register_project(&self, host_user_id: UserId) -> Result<ProjectId> {
-        Ok(sqlx::query_scalar(
-            "
-            INSERT INTO projects(host_user_id)
-            VALUES ($1)
-            RETURNING id
-            ",
-        )
-        .bind(host_user_id)
-        .fetch_one(&self.pool)
-        .await
-        .map(ProjectId)?)
+        test_support!(self, {
+            Ok(sqlx::query_scalar(
+                "
+                INSERT INTO projects(host_user_id)
+                VALUES ($1)
+                RETURNING id
+                ",
+            )
+            .bind(host_user_id)
+            .fetch_one(&self.pool)
+            .await
+            .map(ProjectId)?)
+        })
     }
 
     async fn unregister_project(&self, project_id: ProjectId) -> Result<()> {
-        sqlx::query(
-            "
-            UPDATE projects
-            SET unregistered = 't'
-            WHERE id = $1
-            ",
-        )
-        .bind(project_id)
-        .execute(&self.pool)
-        .await?;
-        Ok(())
+        test_support!(self, {
+            sqlx::query(
+                "
+                UPDATE projects
+                SET unregistered = 't'
+                WHERE id = $1
+                ",
+            )
+            .bind(project_id)
+            .execute(&self.pool)
+            .await?;
+            Ok(())
+        })
     }
 
     async fn update_worktree_extensions(

crates/collab/src/integration_tests.rs 🔗

@@ -53,6 +53,7 @@ use std::{
     time::Duration,
 };
 use theme::ThemeRegistry;
+use tokio::runtime::{EnterGuard, Runtime};
 use unindent::Unindent as _;
 use util::post_inc;
 use workspace::{shared_screen::SharedScreen, Item, SplitDirection, ToggleFollow, Workspace};
@@ -72,8 +73,15 @@ async fn test_basic_calls(
     cx_b2: &mut TestAppContext,
     cx_c: &mut TestAppContext,
 ) {
+    // let runtime = tokio::runtime::Runtime::new().unwrap();
+    // let _enter_guard = runtime.enter();
+
     deterministic.forbid_parking();
     let mut server = TestServer::start(cx_a.foreground(), cx_a.background()).await;
+
+    let start = std::time::Instant::now();
+    eprintln!("test_basic_calls");
+
     let client_a = server.create_client(cx_a, "user_a").await;
     let client_b = server.create_client(cx_b, "user_b").await;
     let client_c = server.create_client(cx_c, "user_c").await;
@@ -259,6 +267,8 @@ async fn test_basic_calls(
             pending: Default::default()
         }
     );
+
+    eprintln!("finished test {:?}", start.elapsed());
 }
 
 #[gpui::test(iterations = 10)]
@@ -6091,7 +6101,12 @@ impl TestServer {
     ) -> Self {
         static NEXT_LIVE_KIT_SERVER_ID: AtomicUsize = AtomicUsize::new(0);
 
-        let test_db = TestDb::fake(background.clone());
+        let test_db = tokio::runtime::Builder::new_current_thread()
+            .enable_io()
+            .enable_time()
+            .build()
+            .unwrap()
+            .block_on(TestDb::postgres());
         let live_kit_server_id = NEXT_LIVE_KIT_SERVER_ID.fetch_add(1, SeqCst);
         let live_kit_server = live_kit_client::TestServer::create(
             format!("http://livekit.{}.test", live_kit_server_id),