Merge pull request #10040 from Turbo87/async-test-builders

tests/builders: Migrate to async/await

Author: Turbo87

src/index.rs

@@ -145,8 +145,7 @@ mod tests {
     #[tokio::test]
     async fn test_index_metadata() {
         let test_db = TestDatabase::new();
-        let mut conn = test_db.connect();
-        let mut async_conn = test_db.async_connect().await;
+        let mut conn = test_db.async_connect().await;
 
         let user_id = diesel::insert_into(users::table)
             .values((
@@ -156,7 +155,7 @@ mod tests {
                 users::gh_access_token.eq("some random token"),
             ))
             .returning(users::id)
-            .get_result::<i32>(&mut async_conn)
+            .get_result::<i32>(&mut conn)
             .await
             .unwrap();
 
@@ -172,9 +171,10 @@ mod tests {
 
         let fooo = CrateBuilder::new("foo", user_id)
             .version(VersionBuilder::new("0.1.0"))
-            .expect_build(&mut conn);
+            .expect_build(&mut conn)
+            .await;
 
-        let metadata = index_metadata(&fooo, &mut async_conn).await.unwrap();
+        let metadata = index_metadata(&fooo, &mut conn).await.unwrap();
         assert_json_snapshot!(metadata);
 
         let bar = CrateBuilder::new("bar", user_id)
@@ -190,9 +190,10 @@ mod tests {
                     .dependency(&fooo, None),
             )
             .version(VersionBuilder::new("1.0.1").checksum("0123456789abcdef"))
-            .expect_build(&mut conn);
+            .expect_build(&mut conn)
+            .await;
 
-        let metadata = index_metadata(&bar, &mut async_conn).await.unwrap();
+        let metadata = index_metadata(&bar, &mut conn).await.unwrap();
         assert_json_snapshot!(metadata);
     }
 }

src/models/category.rs

@@ -3,7 +3,8 @@ use diesel::{
     delete, dsl, insert_into, sql_query, ExpressionMethods, QueryDsl, QueryResult,
     TextExpressionMethods,
 };
-use diesel_async::AsyncPgConnection;
+use diesel_async::scoped_futures::ScopedFutureExt;
+use diesel_async::{AsyncConnection, AsyncPgConnection};
 
 use crate::models::Crate;
 use crate::schema::*;
@@ -46,6 +47,50 @@ impl Category {
         categories::table.filter(filter)
     }
 
+    pub async fn async_update_crate(
+        conn: &mut AsyncPgConnection,
+        crate_id: i32,
+        slugs: &[&str],
+    ) -> QueryResult<Vec<String>> {
+        use diesel_async::RunQueryDsl;
+        conn.transaction(|conn| {
+            async move {
+                let categories: Vec<Category> = categories::table
+                    .filter(categories::slug.eq_any(slugs))
+                    .load(conn)
+                    .await?;
+
+                let invalid_categories = slugs
+                    .iter()
+                    .filter(|s| !categories.iter().any(|c| c.slug == **s))
+                    .map(ToString::to_string)
+                    .collect();
+
+                let crate_categories = categories
+                    .iter()
+                    .map(|c| CrateCategory {
+                        category_id: c.id,
+                        crate_id,
+                    })
+                    .collect::<Vec<_>>();
+
+                delete(crates_categories::table)
+                    .filter(crates_categories::crate_id.eq(crate_id))
+                    .execute(conn)
+                    .await?;
+
+                insert_into(crates_categories::table)
+                    .values(&crate_categories)
+                    .execute(conn)
+                    .await?;
+
+                Ok(invalid_categories)
+            }
+            .scope_boxed()
+        })
+        .await
+    }
+
     pub fn update_crate(
         conn: &mut impl Conn,
         crate_id: i32,

src/models/keyword.rs

@@ -1,5 +1,6 @@
 use chrono::NaiveDateTime;
-use diesel_async::AsyncPgConnection;
+use diesel_async::scoped_futures::ScopedFutureExt;
+use diesel_async::{AsyncConnection, AsyncPgConnection};
 
 use crate::models::Crate;
 use crate::schema::*;
@@ -38,6 +39,31 @@ impl Keyword {
             .await
     }
 
+    pub async fn async_find_or_create_all(
+        conn: &mut AsyncPgConnection,
+        names: &[&str],
+    ) -> QueryResult<Vec<Keyword>> {
+        use diesel_async::RunQueryDsl;
+
+        let lowercase_names: Vec<_> = names.iter().map(|s| s.to_lowercase()).collect();
+
+        let new_keywords: Vec<_> = lowercase_names
+            .iter()
+            .map(|s| keywords::keyword.eq(s))
+            .collect();
+
+        diesel::insert_into(keywords::table)
+            .values(&new_keywords)
+            .on_conflict_do_nothing()
+            .execute(conn)
+            .await?;
+
+        keywords::table
+            .filter(keywords::keyword.eq_any(&lowercase_names))
+            .load(conn)
+            .await
+    }
+
     pub fn find_or_create_all(conn: &mut impl Conn, names: &[&str]) -> QueryResult<Vec<Keyword>> {
         use diesel::RunQueryDsl;
 
@@ -67,6 +93,42 @@ impl Keyword {
             && chars.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-' || c == '+')
     }
 
+    pub async fn async_update_crate(
+        conn: &mut AsyncPgConnection,
+        crate_id: i32,
+        keywords: &[&str],
+    ) -> QueryResult<()> {
+        conn.transaction(|conn| {
+            async move {
+                use diesel_async::RunQueryDsl;
+
+                let keywords = Keyword::async_find_or_create_all(conn, keywords).await?;
+
+                diesel::delete(crates_keywords::table)
+                    .filter(crates_keywords::crate_id.eq(crate_id))
+                    .execute(conn)
+                    .await?;
+
+                let crate_keywords = keywords
+                    .into_iter()
+                    .map(|kw| CrateKeyword {
+                        crate_id,
+                        keyword_id: kw.id,
+                    })
+                    .collect::<Vec<_>>();
+
+                diesel::insert_into(crates_keywords::table)
+                    .values(&crate_keywords)
+                    .execute(conn)
+                    .await?;
+
+                Ok(())
+            }
+            .scope_boxed()
+        })
+        .await
+    }
+
     pub fn update_crate(conn: &mut impl Conn, crate_id: i32, keywords: &[&str]) -> QueryResult<()> {
         conn.transaction(|conn| {
             use diesel::RunQueryDsl;
@@ -99,23 +161,27 @@ mod tests {
     use super::*;
     use crates_io_test_db::TestDatabase;
 
-    #[test]
-    fn dont_associate_with_non_lowercased_keywords() {
-        use diesel::RunQueryDsl;
+    #[tokio::test]
+    #[allow(clippy::iter_next_slice)]
+    async fn dont_associate_with_non_lowercased_keywords() {
+        use diesel_async::RunQueryDsl;
 
         let test_db = TestDatabase::new();
-        let conn = &mut test_db.connect();
+        let mut conn = test_db.async_connect().await;
 
         // The code should be preventing lowercased keywords from existing,
         // but if one happens to sneak in there, don't associate crates with it.
 
         diesel::insert_into(keywords::table)
             .values(keywords::keyword.eq("NO"))
-            .execute(conn)
+            .execute(&mut conn)
+            .await
             .unwrap();
 
-        let associated = Keyword::find_or_create_all(conn, &["no"]).unwrap();
+        let associated = Keyword::async_find_or_create_all(&mut conn, &["no"])
+            .await
+            .unwrap();
         assert_eq!(associated.len(), 1);
-        assert_eq!(associated.first().unwrap().keyword, "no");
+        assert_eq!(associated.iter().next().unwrap().keyword, "no");
     }
 }

src/models/krate.rs

@@ -3,7 +3,8 @@ use diesel::associations::Identifiable;
 use diesel::dsl;
 use diesel::pg::Pg;
 use diesel::sql_types::{Bool, Integer, Text};
-use diesel_async::AsyncPgConnection;
+use diesel_async::scoped_futures::ScopedFutureExt;
+use diesel_async::{AsyncConnection, AsyncPgConnection};
 use secrecy::SecretString;
 use thiserror::Error;
 
@@ -124,6 +125,42 @@ impl<'a> NewCrate<'a> {
             .get_result(conn)
     }
 
+    pub async fn async_create(
+        &self,
+        conn: &mut AsyncPgConnection,
+        user_id: i32,
+    ) -> QueryResult<Crate> {
+        use diesel_async::RunQueryDsl;
+
+        conn.transaction(|conn| {
+            async move {
+                let krate: Crate = diesel::insert_into(crates::table)
+                    .values(self)
+                    .on_conflict_do_nothing()
+                    .returning(Crate::as_returning())
+                    .get_result(conn)
+                    .await?;
+
+                let owner = CrateOwner {
+                    crate_id: krate.id,
+                    owner_id: user_id,
+                    created_by: user_id,
+                    owner_kind: OwnerKind::User,
+                    email_notifications: true,
+                };
+
+                diesel::insert_into(crate_owners::table)
+                    .values(&owner)
+                    .execute(conn)
+                    .await?;
+
+                Ok(krate)
+            }
+            .scope_boxed()
+        })
+        .await
+    }
+
     pub fn create(&self, conn: &mut impl Conn, user_id: i32) -> QueryResult<Crate> {
         use diesel::RunQueryDsl;
 

src/models/team.rs

@@ -47,6 +47,19 @@ pub struct NewTeam<'a> {
 }
 
 impl<'a> NewTeam<'a> {
+    pub async fn async_create_or_update(&self, conn: &mut AsyncPgConnection) -> QueryResult<Team> {
+        use diesel::insert_into;
+        use diesel_async::RunQueryDsl;
+
+        insert_into(teams::table)
+            .values(self)
+            .on_conflict(teams::github_id)
+            .do_update()
+            .set(self)
+            .get_result(conn)
+            .await
+    }
+
     pub fn create_or_update(&self, conn: &mut impl Conn) -> QueryResult<Team> {
         use diesel::insert_into;
         use diesel::RunQueryDsl;

src/models/version.rs

@@ -3,7 +3,8 @@ use std::collections::BTreeMap;
 use bon::Builder;
 use chrono::NaiveDateTime;
 use crates_io_index::features::FeaturesMap;
-use diesel_async::AsyncPgConnection;
+use diesel_async::scoped_futures::ScopedFutureExt;
+use diesel_async::{AsyncConnection, AsyncPgConnection};
 use serde::Deserialize;
 
 use crate::models::{Crate, User};
@@ -117,6 +118,36 @@ impl NewVersion<'_> {
             Ok(version)
         })
     }
+
+    pub async fn async_save(
+        &self,
+        conn: &mut AsyncPgConnection,
+        published_by_email: &str,
+    ) -> QueryResult<Version> {
+        use diesel::insert_into;
+        use diesel_async::RunQueryDsl;
+
+        conn.transaction(|conn| {
+            async move {
+                let version: Version = insert_into(versions::table)
+                    .values(self)
+                    .get_result(conn)
+                    .await?;
+
+                insert_into(versions_published_by::table)
+                    .values((
+                        versions_published_by::version_id.eq(version.id),
+                        versions_published_by::email.eq(published_by_email),
+                    ))
+                    .execute(conn)
+                    .await?;
+
+                Ok(version)
+            }
+            .scope_boxed()
+        })
+        .await
+    }
 }
 
 fn strip_build_metadata(version: &str) -> &str {

src/tests/blocked_routes.rs

@@ -11,11 +11,12 @@ async fn test_non_blocked_download_route() {
         .with_user()
         .await;
 
-    let mut conn = app.db_conn();
+    let mut conn = app.async_db_conn().await;
 
     CrateBuilder::new("foo", user.as_model().id)
         .version(VersionBuilder::new("1.0.0"))
-        .expect_build(&mut conn);
+        .expect_build(&mut conn)
+        .await;
 
     let status = anon
         .get::<()>("/api/v1/crates/foo/1.0.0/download")
@@ -36,11 +37,12 @@ async fn test_blocked_download_route() {
         .with_user()
         .await;
 
-    let mut conn = app.db_conn();
+    let mut conn = app.async_db_conn().await;
 
     CrateBuilder::new("foo", user.as_model().id)
         .version(VersionBuilder::new("1.0.0"))
-        .expect_build(&mut conn);
+        .expect_build(&mut conn)
+        .await;
 
     let status = anon
         .get::<()>("/api/v1/crates/foo/1.0.0/download")