AppFlowy-Cloud/tests/file_test/multiple_part_test.rs

417 lines
11 KiB
Rust

use super::TestBucket;
use crate::collab::util::{generate_random_bytes, generate_random_string};
use app_error::ErrorCode;
use appflowy_cloud::api::file_storage::BlobPathV1;
use aws_sdk_s3::types::CompletedPart;
use bytes::Bytes;
use client_api::ChunkedBytes;
use client_api_test::{generate_unique_registered_user_client, workspace_id_from_client};
use database::file::{BlobKey, BucketClient, ResponseBlob};
use database_entity::file_dto::{
CompleteUploadRequest, CompletedPartRequest, CreateUploadRequest, UploadPartData,
};
use uuid::Uuid;
#[tokio::test]
async fn multiple_part_put_and_get_test() {
let (c1, _user1) = generate_unique_registered_user_client().await;
let workspace_id = workspace_id_from_client(&c1).await;
let parent_dir = workspace_id.clone();
let mime = mime::TEXT_PLAIN_UTF_8;
let text = generate_random_string(8 * 1024 * 1024);
let file_id = Uuid::new_v4().to_string();
let upload = c1
.create_upload(
&workspace_id,
CreateUploadRequest {
file_id: file_id.clone(),
parent_dir: parent_dir.clone(),
content_type: mime.to_string(),
},
)
.await
.unwrap();
let mut chunked_bytes = ChunkedBytes::from_bytes(Bytes::from(text.clone())).unwrap();
assert_eq!(chunked_bytes.offsets.len(), 2);
chunked_bytes.set_chunk_size(5 * 1024 * 1024).unwrap();
let mut completed_parts = Vec::new();
let iter = chunked_bytes.iter().enumerate();
for (index, next) in iter {
let resp = c1
.upload_part(
&workspace_id,
&parent_dir,
&file_id,
&upload.upload_id,
index as i32 + 1,
next.to_vec(),
)
.await
.unwrap();
completed_parts.push(CompletedPartRequest {
e_tag: resp.e_tag,
part_number: resp.part_num,
});
}
assert_eq!(completed_parts.len(), 2);
assert_eq!(completed_parts[0].part_number, 1);
assert_eq!(completed_parts[1].part_number, 2);
let req = CompleteUploadRequest {
file_id: file_id.clone(),
parent_dir: parent_dir.clone(),
upload_id: upload.upload_id,
parts: completed_parts,
};
c1.complete_upload(&workspace_id, req).await.unwrap();
let blob = c1
.get_blob_v1(&workspace_id, &parent_dir, &file_id)
.await
.unwrap()
.1;
let blob_text = String::from_utf8(blob.to_vec()).unwrap();
assert_eq!(blob_text, text);
}
#[tokio::test]
async fn single_part_put_and_get_test() {
// Test with smaller file (single part)
let (c1, _user1) = generate_unique_registered_user_client().await;
let workspace_id = workspace_id_from_client(&c1).await;
let mime = mime::TEXT_PLAIN_UTF_8;
let text = generate_random_string(1024);
let file_id = Uuid::new_v4().to_string();
let upload = c1
.create_upload(
&workspace_id,
CreateUploadRequest {
file_id: file_id.clone(),
parent_dir: workspace_id.clone(),
content_type: mime.to_string(),
},
)
.await
.unwrap();
let chunked_bytes = ChunkedBytes::from_bytes(Bytes::from(text.clone())).unwrap();
assert_eq!(chunked_bytes.offsets.len(), 1);
let mut completed_parts = Vec::new();
let iter = chunked_bytes.iter().enumerate();
for (index, next) in iter {
let resp = c1
.upload_part(
&workspace_id,
&workspace_id,
&file_id,
&upload.upload_id,
index as i32 + 1,
next.to_vec(),
)
.await
.unwrap();
completed_parts.push(CompletedPartRequest {
e_tag: resp.e_tag,
part_number: resp.part_num,
});
}
assert_eq!(completed_parts.len(), 1);
assert_eq!(completed_parts[0].part_number, 1);
let req = CompleteUploadRequest {
file_id: file_id.clone(),
parent_dir: workspace_id.clone(),
upload_id: upload.upload_id,
parts: completed_parts,
};
c1.complete_upload(&workspace_id, req).await.unwrap();
let blob = c1
.get_blob_v1(&workspace_id, &workspace_id, &file_id)
.await
.unwrap()
.1;
let blob_text = String::from_utf8(blob.to_vec()).unwrap();
assert_eq!(blob_text, text);
}
#[tokio::test]
async fn empty_part_upload_test() {
// Test with empty part
let (c1, _user1) = generate_unique_registered_user_client().await;
let workspace_id = workspace_id_from_client(&c1).await;
let mime = mime::TEXT_PLAIN_UTF_8;
let file_id = Uuid::new_v4().to_string();
let upload = c1
.create_upload(
&workspace_id,
CreateUploadRequest {
file_id: file_id.clone(),
parent_dir: workspace_id.clone(),
content_type: mime.to_string(),
},
)
.await
.unwrap();
let result = c1
.upload_part(&workspace_id, "", &file_id, &upload.upload_id, 1, vec![])
.await
.unwrap_err();
assert_eq!(result.code, ErrorCode::InvalidRequest)
}
#[tokio::test]
async fn multiple_part_upload_test() {
let test_bucket = TestBucket::new().await;
// Test with a payload of less than 5MB
let small_file_size = 4 * 1024 * 1024; // 4 MB
let small_blob = generate_random_bytes(small_file_size);
perform_upload_test(&test_bucket, small_blob, small_file_size, "small_file").await;
// Test with a payload of exactly 10MB
let file_size = 10 * 1024 * 1024; // 10 MB
let blob = generate_random_bytes(file_size);
perform_upload_test(&test_bucket, blob, file_size, "large_file").await;
// Test with a payload of exactly 20MB
let file_size = 20 * 1024 * 1024; // 20 MB
let blob = generate_random_bytes(file_size);
perform_upload_test(&test_bucket, blob, file_size, "large_file").await;
}
#[tokio::test]
#[should_panic]
async fn multiple_part_upload_empty_data_test() {
let test_bucket = TestBucket::new().await;
let empty_blob = Vec::new();
perform_upload_test(&test_bucket, empty_blob, 0, "empty_file").await;
}
async fn perform_upload_test(
test_bucket: &TestBucket,
blob: Vec<u8>,
file_size: usize,
description: &str,
) {
let chunk_size = 5 * 1024 * 1024; // 5 MB
let file_id = Uuid::new_v4().to_string();
let workspace_id = Uuid::new_v4();
let parent_dir = workspace_id.to_string();
let req = CreateUploadRequest {
file_id: file_id.clone(),
parent_dir: parent_dir.clone(),
content_type: "text".to_string(),
};
let key = BlobPathV1 {
workspace_id,
parent_dir: parent_dir.clone(),
file_id,
};
let upload = test_bucket.create_upload(key, req).await.unwrap();
let mut chunk_count = (file_size / chunk_size) + 1;
let mut size_of_last_chunk = file_size % chunk_size;
if size_of_last_chunk == 0 {
size_of_last_chunk = chunk_size;
chunk_count -= 1;
}
let mut completed_parts = Vec::new();
for chunk_index in 0..chunk_count {
let start = chunk_index * chunk_size;
let end = start
+ if chunk_index == chunk_count - 1 {
size_of_last_chunk
} else {
chunk_size
};
let chunk = &blob[start..end];
let part_number = (chunk_index + 1) as i32;
let req = UploadPartData {
file_id: upload.file_id.clone(),
upload_id: upload.upload_id.clone(),
part_number,
body: chunk.to_vec(),
};
let key = BlobPathV1 {
workspace_id,
parent_dir: parent_dir.clone(),
file_id: upload.file_id.clone(),
};
let resp = test_bucket.upload_part(&key, req).await.unwrap();
completed_parts.push(
CompletedPart::builder()
.e_tag(resp.e_tag)
.part_number(resp.part_num)
.build(),
);
}
let complete_req = CompleteUploadRequest {
file_id: upload.file_id.clone(),
parent_dir: parent_dir.clone(),
upload_id: upload.upload_id.clone(),
parts: completed_parts
.into_iter()
.map(|p| CompletedPartRequest {
e_tag: p.e_tag().unwrap().to_string(),
part_number: p.part_number.unwrap(),
})
.collect(),
};
let key = BlobPathV1 {
workspace_id,
parent_dir: parent_dir.clone(),
file_id: upload.file_id.clone(),
};
test_bucket
.complete_upload(&key, complete_req)
.await
.unwrap();
// Verify the upload
let object = test_bucket.get_blob(&key.object_key()).await.unwrap();
assert_eq!(object.len(), file_size, "Failed for {}", description);
assert_eq!(object.to_blob(), blob, "Failed for {}", description);
}
#[tokio::test]
async fn invalid_test() {
let (c1, _user1) = generate_unique_registered_user_client().await;
let workspace_id = workspace_id_from_client(&c1).await;
let parent_dir = workspace_id.clone();
let file_id = uuid::Uuid::new_v4().to_string();
let mime = mime::TEXT_PLAIN_UTF_8;
// test invalid create upload request
for request in [
CreateUploadRequest {
file_id: "".to_string(),
parent_dir: parent_dir.clone(),
content_type: mime.to_string(),
},
CreateUploadRequest {
file_id: file_id.clone(),
parent_dir: "".to_string(),
content_type: mime.to_string(),
},
] {
let err = c1.create_upload(&workspace_id, request).await.unwrap_err();
assert_eq!(err.code, ErrorCode::InvalidRequest);
}
// test invalid upload part request
let upload_id = uuid::Uuid::new_v4().to_string();
for request in vec![
// workspace_id, parent_dir, file_id, upload_id, part_number, body
(
"".to_string(),
parent_dir.clone(),
file_id.clone(),
upload_id.clone(),
1,
vec![1, 2, 3],
),
(
workspace_id.clone(),
"".to_string(),
file_id.clone(),
upload_id.clone(),
1,
vec![1, 2, 3],
),
(
workspace_id.clone(),
parent_dir.clone(),
"".to_string(),
upload_id.clone(),
1,
vec![1, 2, 3],
),
] {
let err = c1
.upload_part(
&request.0, &request.1, &request.2, &request.3, request.4, request.5,
)
.await
.unwrap_err();
assert_eq!(err.code, ErrorCode::Internal);
}
}
#[tokio::test]
async fn multiple_level_dir_upload_file_test() {
// Test with smaller file (single part)
let (c1, _user1) = generate_unique_registered_user_client().await;
let workspace_id = workspace_id_from_client(&c1).await;
let mime = mime::TEXT_PLAIN_UTF_8;
let text = generate_random_string(1024);
let file_id = Uuid::new_v4().to_string();
let parent_dir = "file/v1/image".to_string();
let upload = c1
.create_upload(
&workspace_id,
CreateUploadRequest {
file_id: file_id.clone(),
parent_dir: parent_dir.clone(),
content_type: mime.to_string(),
},
)
.await
.unwrap();
let chunked_bytes = ChunkedBytes::from_bytes(Bytes::from(text.clone())).unwrap();
let mut completed_parts = Vec::new();
let iter = chunked_bytes.iter().enumerate();
for (index, next) in iter {
let resp = c1
.upload_part(
&workspace_id,
&parent_dir,
&file_id,
&upload.upload_id,
index as i32 + 1,
next.to_vec(),
)
.await
.unwrap();
completed_parts.push(CompletedPartRequest {
e_tag: resp.e_tag,
part_number: resp.part_num,
});
}
let req = CompleteUploadRequest {
file_id: file_id.clone(),
parent_dir: parent_dir.clone(),
upload_id: upload.upload_id,
parts: completed_parts,
};
c1.complete_upload(&workspace_id, req).await.unwrap();
let blob = c1
.get_blob_v1(&workspace_id, &parent_dir, &file_id)
.await
.unwrap()
.1;
let blob_text = String::from_utf8(blob.to_vec()).unwrap();
assert_eq!(blob_text, text);
}