From 130e0daa084ddb2c5068c95ffdefadcfd83c85d7 Mon Sep 17 00:00:00 2001 From: Michael Chung Date: Mon, 4 Dec 2023 22:49:33 -0500 Subject: [PATCH] paging works correctly. --- .../src/queries/model_weights/tests.rs | 389 +++++++++++++----- 1 file changed, 290 insertions(+), 99 deletions(-) diff --git a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs index 2a24c7fda3..917fdd98b3 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs +++ b/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(()) + } +}