builders/krate: Add category option

Author: Turbo87

src/tests/builders/krate.rs

@@ -1,5 +1,5 @@
 use cargo_registry::{
-    models::{Crate, Keyword, NewCrate},
+    models::{Category, Crate, Keyword, NewCrate},
     schema::{crates, version_downloads},
     util::errors::AppResult,
 };
@@ -13,6 +13,7 @@ use super::VersionBuilder;
 /// If you want to test logic that happens as part of a publish request, use `PublishBuilder`
 /// instead.
 pub struct CrateBuilder<'a> {
+    categories: Vec<&'a str>,
     downloads: Option<i32>,
     keywords: Vec<&'a str>,
     krate: NewCrate<'a>,
@@ -27,6 +28,7 @@ impl<'a> CrateBuilder<'a> {
     /// doesn't exist in the database, `expect_build` will fail.
     pub fn new(name: &str, owner_id: i32) -> CrateBuilder<'_> {
         CrateBuilder {
+            categories: Vec::new(),
             downloads: None,
             keywords: Vec::new(),
             krate: NewCrate {
@@ -91,6 +93,12 @@ impl<'a> CrateBuilder<'a> {
         self
     }
 
+    /// Adds a category to the crate.
+    pub fn category(mut self, category: &'a str) -> Self {
+        self.categories.push(category);
+        self
+    }
+
     /// Adds a keyword to the crate.
     pub fn keyword(mut self, keyword: &'a str) -> Self {
         self.keywords.push(keyword);
@@ -143,6 +151,10 @@ impl<'a> CrateBuilder<'a> {
             select(refresh_recent_crate_downloads).execute(connection)?;
         }
 
+        if !self.categories.is_empty() {
+            Category::update_crate(connection, &krate, &self.categories)?;
+        }
+
         if !self.keywords.is_empty() {
             Keyword::update_crate(connection, &krate, &self.keywords)?;
         }

src/tests/krate.rs

@@ -1341,17 +1341,23 @@ fn summary_new_crates() {
         let now_ = Utc::now().naive_utc();
         let now_plus_two = now_ + chrono::Duration::seconds(2);
 
-        let krate = CrateBuilder::new("some_downloads", user.id)
+        new_category("Category 1", "cat1", "Category 1 crates")
+            .create_or_update(conn)
+            .unwrap();
+
+        CrateBuilder::new("some_downloads", user.id)
             .version(VersionBuilder::new("0.1.0"))
             .description("description")
             .keyword("popular")
+            .category("cat1")
             .downloads(20)
             .recent_downloads(10)
             .expect_build(conn);
 
-        let krate2 = CrateBuilder::new("most_recent_downloads", user.id)
+        CrateBuilder::new("most_recent_downloads", user.id)
             .version(VersionBuilder::new("0.2.0"))
             .keyword("popular")
+            .category("cat1")
             .downloads(5000)
             .recent_downloads(50)
             .expect_build(conn);
@@ -1377,12 +1383,6 @@ fn summary_new_crates() {
             .downloads(1000)
             .expect_build(conn);
 
-        new_category("Category 1", "cat1", "Category 1 crates")
-            .create_or_update(conn)
-            .unwrap();
-        Category::update_crate(conn, &krate, &["cat1"]).unwrap();
-        Category::update_crate(conn, &krate2, &["cat1"]).unwrap();
-
         // set total_downloads global value for `num_downloads` prop
         update(metadata::table)
             .set(metadata::total_downloads.eq(6000))