paging works correctly.

crates/lib/mysql_queries/src/queries/model_weights/tests.rs

@@ -3,8 +3,9 @@ mod tests {
     use std::result;
 
     use anyhow::Ok;
-    use enums::by_table::model_weights::{weights_types, weights_category};
-    use sqlx::{MySqlPool, Acquire};
+    use chrono::offset;
+    use enums::by_table::model_weights::{ weights_types, weights_category };
+    use sqlx::{ MySqlPool, Acquire };
     use rand::Rng;
 
     use tokio;
@@ -24,7 +25,12 @@ mod tests {
     use crate::queries::model_weights::create_weight::create_weight;
     use crate::queries::model_weights::get_weight::get_weight_by_token;
 
-    use crate::queries::model_weights::delete_weights::{ delete_weights_as_user, delete_weights_as_mod, undelete_weights_as_mod , undelete_weights_as_user};
+    use crate::queries::model_weights::delete_weights::{
+        delete_weights_as_user,
+        delete_weights_as_mod,
+        undelete_weights_as_mod,
+        undelete_weights_as_user,
+    };
     use crate::queries::model_weights::list_weights_by_user::list_weights_by_creator_username;
 
     use crate::queries::model_weights::list_weights_query_builder::ListWeightsQueryBuilder;
@@ -45,23 +51,19 @@ mod tests {
 
     pub async fn delete_all_weights_for_table(mysql_pool: &MySqlPool) -> AnyhowResult<()> {
         // write a query that deletes all weights
-        let _r = sqlx::query("DELETE FROM model_weights")
-        .execute(mysql_pool)
-        .await?;
+        let _r = sqlx::query("DELETE FROM model_weights").execute(mysql_pool).await?;
         Ok(())
     }
 
     #[tokio::test]
     #[serial]
     async fn test_create_weights() -> AnyhowResult<()> {
-
         let pool = setup().await;
         // create a random token for the model weight
         let mut rng = rand::thread_rng();
         let random_number: u32 = rng.gen();
         let model_weight_token1 = ModelWeightToken(random_number.to_string());
 
-
         let creator_token1 = UserToken("creatorToken!1".to_string());
 
         let args = CreateModelWeightsArgs {
@@ -131,7 +133,6 @@ mod tests {
     #[tokio::test]
     #[serial]
     async fn test_delete_and_undelete_weights_user() -> AnyhowResult<()> {
-
         let pool = setup().await;
         // create a random token for the model weight
         let mut rng = rand::thread_rng();
@@ -170,7 +171,7 @@ mod tests {
         create_weight(args).await?;
 
         delete_weights_as_user(&model_weight_token1, &pool).await?;
-        
+
         let result = get_weight_by_token(&model_weight_token1, true, &pool).await?;
         let result = result.unwrap();
 
@@ -261,7 +262,6 @@ mod tests {
             }
         }
 
-
         undelete_weights_as_mod(&model_weight_token1, &pool).await?;
         let result = get_weight_by_token(&model_weight_token1, true, &pool).await?;
         let result = result.unwrap();
@@ -285,13 +285,22 @@ mod tests {
     #[serial]
     async fn list_all_weights_by_user() -> AnyhowResult<()> {
         let pool = setup().await;
-        
+
         let creator_username = "hanashi".to_string();
         let can_see_deleted = true;
 
-        seed_weights_with_category_and_type(WeightsType::LoRA,WeightsCategory::ImageGeneration,5,&creator_username).await?;
+        seed_weights_with_category_and_type(
+            WeightsType::LoRA,
+            WeightsCategory::ImageGeneration,
+            5,
+            &creator_username
+        ).await?;
 
-        let weights_by_username = list_weights_by_creator_username(&pool, &creator_username, can_see_deleted).await?;
+        let weights_by_username = list_weights_by_creator_username(
+            &pool,
+            &creator_username,
+            can_see_deleted
+        ).await?;
         for weight in weights_by_username.iter() {
             // print weight
             println!("weight: {:?}", weight.creator_username);
@@ -301,23 +310,29 @@ mod tests {
         Ok(())
     }
 
-    async fn seed_weights_with_category_and_type(weights_type:WeightsType,weights_category:WeightsCategory,number_of_items:u32,creator_username:&str) -> AnyhowResult<()> {
-
+    async fn seed_weights_with_category_and_type(
+        weights_type: WeightsType,
+        weights_category: WeightsCategory,
+        number_of_items: u32,
+        creator_username: &str
+    ) -> AnyhowResult<()> {
         let pool = setup().await;
-        let creator_token = get_user_token_by_username(&creator_username, &pool).await?.unwrap_or_else(
-            || panic!("Could not find user with username {}", creator_username)
+        let creator_token = get_user_token_by_username(
+            &creator_username,
+            &pool
+        ).await?.unwrap_or_else(||
+            panic!("Could not find user with username {}", creator_username)
         );
 
         for i in 0..number_of_items {
-
             let mut rng = rand::thread_rng();
             let number: u32 = rng.gen();
 
             let model_weight_token = ModelWeightToken(number.to_string());
-                 
+
             let args = CreateModelWeightsArgs {
                 token: &model_weight_token, // replace with actual ModelWeightToken
-                weights_type:weights_type, // replace with actual WeightsType
+                weights_type: weights_type, // replace with actual WeightsType
                 weights_category: weights_category, // replace with actual WeightsCategory
                 title: format!("Title {}", i),
                 maybe_thumbnail_token: Some(format!("Thumbnail {}", i)),
@@ -346,7 +361,6 @@ mod tests {
         }
 
         Ok(())
-
     }
 
     // Tests paging for the list weights query builder
@@ -356,40 +370,45 @@ mod tests {
         let pool = setup().await;
         let creator_username = "hanashi".to_string();
 
-        seed_weights_with_category_and_type(WeightsType::LoRA,WeightsCategory::ImageGeneration,20,&creator_username).await?;
-        
+        seed_weights_with_category_and_type(
+            WeightsType::LoRA,
+            WeightsCategory::ImageGeneration,
+            20,
+            &creator_username
+        ).await?;
+
         let qb = ListWeightsQueryBuilder::new()
-        .offset(Some(0))
-        .weights_type(Some(WeightsType::LoRA))
-        .weights_category(Some(WeightsCategory::ImageGeneration))
-        .limit(10)
-        .scope_creator_username(Some("hanashi"))
-        .include_user_deleted_results(false);
-       
+            .offset(Some(0))
+            .weights_type(Some(WeightsType::LoRA))
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(10)
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false);
+
         let result = qb.perform_query_for_page(&pool).await?;
-    
+
         // let result = qb.perform_query_for_page(&pool).await?;
         assert_eq!(result.weights.len(), 10);
 
         let qb = ListWeightsQueryBuilder::new()
-        .weights_type(Some(WeightsType::LoRA))
-        .weights_category(Some(WeightsCategory::ImageGeneration))
-        .scope_creator_username(Some("hanashi"))
-        .include_user_deleted_results(false)
-        .offset(Some(10))
-        .limit(10);
+            .weights_type(Some(WeightsType::LoRA))
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .offset(Some(10))
+            .limit(10);
 
         let result = qb.perform_query_for_page(&pool).await?;
 
         assert_eq!(result.weights.len(), 10);
 
         let qb = ListWeightsQueryBuilder::new()
-        .weights_type(Some(WeightsType::LoRA))
-        .weights_category(Some(WeightsCategory::ImageGeneration))
-        .scope_creator_username(Some("hanashi"))
-        .include_user_deleted_results(false)
-        .offset(Some(10))
-        .limit(10);
+            .weights_type(Some(WeightsType::LoRA))
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .offset(Some(10))
+            .limit(10);
 
         let result = qb.perform_query_for_page(&pool).await?;
 
@@ -398,27 +417,31 @@ mod tests {
         Ok(())
     }
 
-
     #[tokio::test]
     #[serial]
     async fn list_weights_query_build_test_asc_desc_cursor_reverse() -> AnyhowResult<()> {
         let pool = setup().await;
         let creator_username = "hanashi".to_string();
 
-        seed_weights_with_category_and_type(WeightsType::LoRA,WeightsCategory::ImageGeneration,10,&creator_username).await?;
-        
+        seed_weights_with_category_and_type(
+            WeightsType::LoRA,
+            WeightsCategory::ImageGeneration,
+            10,
+            &creator_username
+        ).await?;
+
         let qb = ListWeightsQueryBuilder::new()
-        .offset(Some(0))
-        .weights_type(Some(WeightsType::LoRA))
-        .weights_category(Some(WeightsCategory::ImageGeneration))
-        .limit(10)
-        .scope_creator_username(Some("hanashi"))
-        .include_user_deleted_results(false)
-        .sort_ascending(true)
-        .cursor_is_reversed(false);
+            .offset(Some(0))
+            .weights_type(Some(WeightsType::LoRA))
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(10)
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .sort_ascending(true)
+            .cursor_is_reversed(false);
 
         let result = qb.perform_query_for_page(&pool).await?;
-        
+
         // write code that will loop through the result and check if the ids are in ascending order
         let mut previous_id = 0;
         for weight in result.weights.iter() {
@@ -428,19 +451,18 @@ mod tests {
         }
 
         let qb = ListWeightsQueryBuilder::new()
-        .offset(Some(0))
-        .weights_type(Some(WeightsType::LoRA))
-        .weights_category(Some(WeightsCategory::ImageGeneration))
-        .limit(10)
-        .scope_creator_username(Some("hanashi"))
-        .include_user_deleted_results(false)
-        .sort_ascending(false)
-        .cursor_is_reversed(true);
+            .offset(Some(0))
+            .weights_type(Some(WeightsType::LoRA))
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(10)
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .sort_ascending(false)
+            .cursor_is_reversed(true);
         let result = qb.perform_query_for_page(&pool).await?;
 
         // write code that will loop through the result and check if the ids are in descending order
-        let mut previous_id
-        = 10000;
+        let mut previous_id = 10000;
         for weight in result.weights.iter() {
             let current_id = weight.weight_id;
             assert!(current_id < previous_id);
@@ -456,18 +478,23 @@ mod tests {
         let pool = setup().await;
         let creator_username = "hanashi".to_string();
 
-        seed_weights_with_category_and_type(WeightsType::LoRA,WeightsCategory::ImageGeneration,10,&creator_username).await?;
-        
+        seed_weights_with_category_and_type(
+            WeightsType::LoRA,
+            WeightsCategory::ImageGeneration,
+            10,
+            &creator_username
+        ).await?;
+
         let qb = ListWeightsQueryBuilder::new()
-        .offset(Some(0))
-        .weights_type(Some(WeightsType::LoRA))
-        .weights_category(Some(WeightsCategory::ImageGeneration))
-        .limit(10)
-        .scope_creator_username(Some("hanashi"))
-        .include_user_deleted_results(false)
-        .sort_ascending(true);
+            .offset(Some(0))
+            .weights_type(Some(WeightsType::LoRA))
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(10)
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .sort_ascending(true);
         let result = qb.perform_query_for_page(&pool).await?;
-        
+
         // write code that will loop through the result and check if the ids are in ascending order
         let mut previous_id = 0;
         for weight in result.weights.iter() {
@@ -477,19 +504,18 @@ mod tests {
         }
 
         let qb = ListWeightsQueryBuilder::new()
-        .offset(Some(0))
-        .weights_type(Some(WeightsType::LoRA))
-        .weights_category(Some(WeightsCategory::ImageGeneration))
-        .limit(10)
-        .scope_creator_username(Some("hanashi"))
-        .include_user_deleted_results(false)
-        .sort_ascending(false);
+            .offset(Some(0))
+            .weights_type(Some(WeightsType::LoRA))
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(10)
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .sort_ascending(false);
         let result = qb.perform_query_for_page(&pool).await?;
 
         // write code that will loop through the result and check if the ids are in descending order
 
-        let mut previous_id
-        = 10000;
+        let mut previous_id = 10000;
         for weight in result.weights.iter() {
             let current_id = weight.weight_id;
             assert!(current_id < previous_id);
@@ -503,11 +529,14 @@ mod tests {
     #[serial]
     async fn list_weights_query_build_test() -> AnyhowResult<()> {
         let pool = setup().await;
-        
+
         let creator_username = "hanashi".to_string();
 
-        let creator_token = get_user_token_by_username(&creator_username, &pool).await?.unwrap_or_else(
-            || panic!("Could not find user with username {}", creator_username)
+        let creator_token = get_user_token_by_username(
+            &creator_username,
+            &pool
+        ).await?.unwrap_or_else(||
+            panic!("Could not find user with username {}", creator_username)
         );
 
         for i in 0..6 {
@@ -519,7 +548,7 @@ mod tests {
                 WeightsType::RvcV2
             } else if i % 5 == 3 {
                 WeightsType::StableDiffusionXL
-            } else if i ^ 5 == 4 {
+            } else if (i ^ 5) == 4 {
                 WeightsType::SoVitsSvc
             } else {
                 WeightsType::StableDiffusion15
@@ -533,7 +562,7 @@ mod tests {
                 WeightsCategory::VoiceConversion
             } else if i % 5 == 3 {
                 WeightsCategory::ImageGeneration
-            } else if i ^ 5 == 4 {
+            } else if (i ^ 5) == 4 {
                 WeightsCategory::VoiceConversion
             } else {
                 WeightsCategory::ImageGeneration
@@ -543,10 +572,10 @@ mod tests {
             let number: u32 = rng.gen();
 
             let model_weight_token = ModelWeightToken(number.to_string());
-                 
+
             let args = CreateModelWeightsArgs {
                 token: &model_weight_token, // replace with actual ModelWeightToken
-                weights_type:weights_type, // replace with actual WeightsType
+                weights_type: weights_type, // replace with actual WeightsType
                 weights_category: weights_category, // replace with actual WeightsCategory
                 title: format!("Title {}", i),
                 maybe_thumbnail_token: Some(format!("Thumbnail {}", i)),
@@ -581,18 +610,18 @@ mod tests {
             .scope_creator_username(Some("hanashi"))
             .include_user_deleted_results(false);
 
-
         let result = qb.perform_query_for_page(&pool).await?;
 
         assert_eq!(result.weights.len(), 2);
 
         let qb = ListWeightsQueryBuilder::new()
-        .weights_type(Some(WeightsType::LoRA))
-        .weights_category(Some(WeightsCategory::ImageGeneration))
-        .scope_creator_username(Some("hanashi"))
-        .include_user_deleted_results(false);
+            .weights_type(Some(WeightsType::LoRA))
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false);
 
         let result = qb.perform_query_for_page(&pool).await?;
+
         assert_eq!(result.weights.len(), 1);
         assert_eq!(result.weights[0].weights_category, WeightsCategory::ImageGeneration);
         assert_eq!(result.weights[0].weights_type, WeightsType::LoRA);
@@ -600,6 +629,168 @@ mod tests {
         Ok(())
     }
 
-    
-}
+    #[tokio::test]
+    #[serial]
+    async fn list_weights_query_builder_page_type_and_category() -> AnyhowResult<()> {
+        let db_connection_string = DEFAULT_MYSQL_CONNECTION_STRING;
+        let pool = MySqlPoolOptions::new()
+            .max_connections(3)
+            .connect(&db_connection_string).await
+            .unwrap();
+        // using the seeded weights from the previous test
+        let creator_username = "hanashi".to_string();
 
+        for i in 0..4 {
+            let weights_types = vec![
+                WeightsType::HifiganTacotron2,
+                WeightsType::LoRA,
+                WeightsType::RvcV2,
+                WeightsType::StableDiffusionXL
+            ];
+
+            let weights_type = weights_types[i as usize];
+
+            let qb = ListWeightsQueryBuilder::new()
+                .weights_type(Some(weights_type))
+                .scope_creator_username(Some("hanashi"))
+                .include_user_deleted_results(false);
+
+            println!("weights_type: {:?}", weights_type);
+            let result = qb.perform_query_for_page(&pool).await?;
+            assert_eq!(result.weights.len(), 20, "Should be 20 weights");
+            assert_eq!(
+                result.weights[0].weights_type,
+                weights_type,
+                "Should be the same weights type"
+            );
+        }
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn list_weight_query_builder_all() -> AnyhowResult<()> {
+        let db_connection_string = DEFAULT_MYSQL_CONNECTION_STRING;
+        let pool = MySqlPoolOptions::new()
+            .max_connections(3)
+            .connect(&db_connection_string).await
+            .unwrap();
+
+        let qb = ListWeightsQueryBuilder::new()
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .limit(11);
+
+        let result = qb.perform_query_for_page(&pool).await?;
+        assert_eq!(result.weights.len(), 11);
+        Ok(())
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn list_limits() -> AnyhowResult<()> {
+        let db_connection_string = DEFAULT_MYSQL_CONNECTION_STRING;
+
+        let pool = MySqlPoolOptions::new()
+            .max_connections(3)
+            .connect(&db_connection_string).await
+            .unwrap();
+
+        let qb = ListWeightsQueryBuilder::new()
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(10);
+
+        let result = qb.perform_query_for_page(&pool).await?;
+
+        assert_eq!(result.weights.len(), 10);
+        assert_eq!(result.weights[0].weights_category, WeightsCategory::ImageGeneration);
+
+        let qb = ListWeightsQueryBuilder::new()
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false);
+
+        let result = qb.perform_query_for_page(&pool).await?;
+
+        // by default 25
+        assert_eq!(result.weights.len(), 25);
+        assert_eq!(result.weights[0].weights_category, WeightsCategory::ImageGeneration);
+        // over query limit
+        let qb: ListWeightsQueryBuilder = ListWeightsQueryBuilder::new()
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(80);
+        let result = qb.perform_query_for_page(&pool).await?;
+        assert_eq!(result.weights.len(), 60);
+        // pick a random value in here
+        assert_eq!(result.weights[9].weights_category, WeightsCategory::ImageGeneration);
+        Ok(())
+    }
+
+    #[tokio::test]
+    #[serial]
+    async fn page_across_a_category() -> AnyhowResult<()> {
+        // paging across a category and it works
+        let db_connection_string = DEFAULT_MYSQL_CONNECTION_STRING;
+        // how we implement paging.
+        let page:u16 = 0;
+        let page_size: u16 = 20;
+
+        let pool = MySqlPoolOptions::new()
+            .max_connections(3)
+            .connect(&db_connection_string).await
+            .unwrap();
+
+        let qb = ListWeightsQueryBuilder::new()
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .sort_ascending(true)
+            .limit(page_size);
+
+        let result = qb.perform_query_for_page(&pool).await?;
+        assert_eq!(result.weights.len(), page_size as usize);
+        assert_eq!(result.weights[0].weights_category, WeightsCategory::ImageGeneration);
+
+        let page = 1;
+        let offset = page * page_size;
+
+        assert_eq!(result.weights.len(), page_size as usize);
+        assert_eq!(result.weights[0].weights_category, WeightsCategory::ImageGeneration);
+        
+
+        let qb = ListWeightsQueryBuilder::new()
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(page_size)
+            .sort_ascending(true)
+            .offset(Some(offset.into()));
+        
+        let result = qb.perform_query_for_page(&pool).await?;
+
+        let page = 2;
+        let offset = page * page_size;
+
+        assert_eq!(result.weights.len(), page_size as usize);
+        assert_eq!(result.weights[0].weights_category, WeightsCategory::ImageGeneration);
+
+        let qb = ListWeightsQueryBuilder::new()
+            .scope_creator_username(Some("hanashi"))
+            .include_user_deleted_results(false)
+            .sort_ascending(true)
+            .weights_category(Some(WeightsCategory::ImageGeneration))
+            .limit(page_size)
+            .offset(Some(offset.into()));
+
+        let result = qb.perform_query_for_page(&pool).await?;
+
+        assert_eq!(result.weights.len(), page_size as usize);
+        assert_eq!(result.weights[0].weights_category, WeightsCategory::ImageGeneration);
+
+        Ok(())
+    }
+}