From d8db78bf7ff297978eda65de25f6afda7f25639a Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 7 Aug 2024 01:57:06 -0400 Subject: [PATCH] tests for gptsovits archives --- Cargo.lock | 1 + crates/service/job/inference_job/Cargo.toml | 3 + .../gpt_sovits/model_package/model_package.rs | 2 +- .../extract_and_upload_gpt_sovits_package.rs | 85 +++++++++++++++++- test_data/archive/zip/empty.zip | Bin 0 -> 172 bytes test_data/archive/zip/gptsovits/all_files.zip | Bin 0 -> 974 bytes .../archive/zip/gptsovits/just_model.zip | Bin 0 -> 576 bytes 7 files changed, 86 insertions(+), 5 deletions(-) create mode 100644 test_data/archive/zip/empty.zip create mode 100644 test_data/archive/zip/gptsovits/all_files.zip create mode 100644 test_data/archive/zip/gptsovits/just_model.zip diff --git a/Cargo.lock b/Cargo.lock index 62e0dd253b..e40ef0a406 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3618,6 +3618,7 @@ dependencies = [ "subprocess", "subprocess_common", "tempdir", + "testing", "thumbnail_generator", "time", "tokens", diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 960bc35e1c..56a4e57542 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -96,3 +96,6 @@ opentelemetry-otlp = { version = "0.16.0", features = ["metrics"] } opentelemetry_sdk = { version = "0.23.0", features = ["metrics", "rt-tokio"] } zip = "2.1.3" url = "2.5.1" + +[dev-dependencies] +testing = { path = "../../../lib/testing" } diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/model_package.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/model_package.rs index ff3ae32d13..0982e9300e 100644 --- a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/model_package.rs +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/model_package/model_package.rs @@ -88,7 +88,7 @@ impl GptSovitsPackageFileType { } } -#[derive(Debug)] +#[derive(Debug, Eq, PartialEq)] pub enum GptSovitsPackageError { InvalidArchive, InvalidGPTModel(String), diff --git a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_upload_gpt_sovits_package.rs b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_upload_gpt_sovits_package.rs index 3276f06dc8..4d28def18f 100644 --- a/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_upload_gpt_sovits_package.rs +++ b/crates/service/job/inference_job/src/job/job_types/gpt_sovits/upload_model/extract_and_upload_gpt_sovits_package.rs @@ -1,4 +1,4 @@ -use std::io::{BufReader, Cursor, Read}; +use std::io::{BufReader, Cursor, Read, Seek}; use std::path::PathBuf; use log::{error, info, warn}; @@ -188,7 +188,7 @@ pub async fn extract_and_upload_gpt_sovits_package_files( } #[derive(Debug,Clone)] -struct PackageZipEntryDetails { +pub struct PackageZipEntryDetails { package_type: GptSovitsPackageFileType, enclosed_name: PathBuf, maybe_better_alternative_output_name: String, @@ -196,7 +196,7 @@ struct PackageZipEntryDetails { is_valid_file_extension: bool, } -fn get_relevant_zip_entries(archive: &mut ZipArchive>>) -> Result, GptSovitsPackageError> { +fn get_relevant_zip_entries(archive: &mut ZipArchive>) -> Result, GptSovitsPackageError> { let mut entries: Vec = Vec::new(); for i in 0..archive.len() { @@ -269,7 +269,16 @@ fn get_relevant_zip_entries(archive: &mut ZipArchive>>) } } - if entries.len() < GptSovitsPackageFileType::all_variants().len() { + // Only the model files are truly essential + let has_no_gpt = entries.iter() + .find(|entry| entry.package_type == GptSovitsPackageFileType::GptModel) + .is_none(); + + let has_no_sovits = entries.iter() + .find(|entry| entry.package_type == GptSovitsPackageFileType::SovitsCheckpoint) + .is_none(); + + if has_no_gpt || has_no_sovits { return Err(GptSovitsPackageError::InvalidArchive); } @@ -279,3 +288,71 @@ fn get_relevant_zip_entries(archive: &mut ZipArchive>>) Ok(entries) } + +#[cfg(test)] +mod tests { + use std::fs::File; + use std::io::BufReader; + + use zip::ZipArchive; + + use testing::test_file_path::test_file_path; + + use crate::job::job_types::gpt_sovits::model_package::model_package::{GptSovitsPackageError, GptSovitsPackageFileType}; + use crate::job::job_types::gpt_sovits::upload_model::extract_and_upload_gpt_sovits_package::{get_relevant_zip_entries, PackageZipEntryDetails}; + + fn read_archive(path: &str) -> ZipArchive> { + let path = test_file_path(path).unwrap(); + let file = File::open(path).unwrap(); + let reader = BufReader::new(file); + ZipArchive::new(reader).unwrap() + } + + #[test] + fn test_empty_archive() { + let mut archive = read_archive("test_data/archive/zip/empty.zip"); + let result = get_relevant_zip_entries(&mut archive); + assert!(result.is_err()); + assert_eq!(result.err().unwrap(), GptSovitsPackageError::InvalidArchive); + } + + #[test] + fn test_archive_with_just_model() { + let mut archive = read_archive("test_data/archive/zip/gptsovits/just_model.zip"); + let result = get_relevant_zip_entries(&mut archive); + + assert!(result.is_ok()); + + let result = result.unwrap(); + assert_eq!(result.len(), 2); + + let entry = result.iter().find(|entry| entry.package_type == GptSovitsPackageFileType::GptModel).unwrap(); + assert_eq!(entry.enclosed_name.to_str().unwrap(), "just_model/gpt.ckpt"); + + let entry = result.iter().find(|entry| entry.package_type == GptSovitsPackageFileType::SovitsCheckpoint).unwrap(); + assert_eq!(entry.enclosed_name.to_str().unwrap(), "just_model/sovits.pth"); + } + + #[test] + fn test_archive_with_single_files() { + let mut archive = read_archive("test_data/archive/zip/gptsovits/all_files.zip"); + let result = get_relevant_zip_entries(&mut archive); + + assert!(result.is_ok()); + + let result = result.unwrap(); + assert_eq!(result.len(), 4); + + let entry = result.iter().find(|entry| entry.package_type == GptSovitsPackageFileType::GptModel).unwrap(); + assert_eq!(entry.enclosed_name.to_str().unwrap(), "all_files/gpt.ckpt"); + + let entry = result.iter().find(|entry| entry.package_type == GptSovitsPackageFileType::SovitsCheckpoint).unwrap(); + assert_eq!(entry.enclosed_name.to_str().unwrap(), "all_files/sovits.pth"); + + let entry = result.iter().find(|entry| entry.package_type == GptSovitsPackageFileType::ReferenceAudio).unwrap(); + assert_eq!(entry.enclosed_name.to_str().unwrap(), "all_files/reference.wav"); + + let entry = result.iter().find(|entry| entry.package_type == GptSovitsPackageFileType::ReferenceTranscript).unwrap(); + assert_eq!(entry.enclosed_name.to_str().unwrap(), "all_files/reference.txt"); + } +} \ No newline at end of file diff --git a/test_data/archive/zip/empty.zip b/test_data/archive/zip/empty.zip new file mode 100644 index 0000000000000000000000000000000000000000..6410163e2bf7f2a5fd69e21f804298e089d17c63 GIT binary patch literal 172 zcmWIWW@Zs#0D+Ah?2%vwlwf60U?@o~F3}GS;bmZ-$FMmKgy#csX$3a}BgAU&*#PZ6grh;M$y4TEWf0 z$nuqufdQ;P0B8;u0|yp!#8J#CN=-{GO3h17)hnqefg1`k7=$qlWe1rE2M9|ZqB&Lw z#gg=b620W?0*Lw3f#$+6(dNHIGar{H$`i}Lo>&I71cr&WDb_0} z$pG6h2WSHf6SRSmNrV}9tir>jL3n5@B~qDj<5t)8nZgQ$Pk?4z0b&LQ0Eslc AlmGw# literal 0 HcmV?d00001 diff --git a/test_data/archive/zip/gptsovits/just_model.zip b/test_data/archive/zip/gptsovits/just_model.zip new file mode 100644 index 0000000000000000000000000000000000000000..723ea2a88a624dd84f9262bbd2a277a2b580e174 GIT binary patch literal 576 zcmWIWW@Zs#00ApL_DC=TN^mnMFl3b$m&E7hr=;fShlcPnurqRRPGbXN5M5fq&A`a= zm63q~Y(N0eAT9fqB;0c-q zqNo-W=a*%c6zdg~WPq)h12i6liL!!`NrV}9aKPO1))B-aDR3B>xR`MV56s{tjh?6m j6CYRz3$O(b!h%((7LXEHtZX1-n1N6h80)-1%)kHu0R3rI literal 0 HcmV?d00001