423 lines
11 KiB
Rust
423 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.object_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.object_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.object_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);
|
|
}
|