use crate::common::api_endpoints::DriveApi;
use openlark_core::{
SDKResult, api::Response, config::Config, http::Transport, req_option::RequestOption,
validate_required,
};
const DEFAULT_MAX_DOWNLOAD_SIZE: usize = 100 * 1024 * 1024;
#[derive(Debug)]
pub struct DownloadMediaRequest {
config: Config,
pub file_token: String,
pub extra: Option<String>,
pub range: Option<String>,
max_size: usize,
}
impl DownloadMediaRequest {
pub fn new(config: Config, file_token: impl Into<String>) -> Self {
Self {
config,
file_token: file_token.into(),
extra: None,
range: None,
max_size: DEFAULT_MAX_DOWNLOAD_SIZE,
}
}
pub fn extra(mut self, extra: impl Into<String>) -> Self {
self.extra = Some(extra.into());
self
}
pub fn range(mut self, range: impl Into<String>) -> Self {
self.range = Some(range.into());
self
}
pub fn max_size(mut self, max_size: usize) -> Self {
self.max_size = max_size;
self
}
pub async fn execute(self) -> SDKResult<Response<Vec<u8>>> {
self.execute_with_options(RequestOption::default()).await
}
pub async fn execute_with_options(self, option: RequestOption) -> SDKResult<Response<Vec<u8>>> {
validate_required!(self.file_token, "file_token 不能为空");
if let Some(range) = &self.range
&& (!range.starts_with("bytes=") || !range.contains('-'))
{
return Err(openlark_core::error::validation_error(
"range",
"range 格式必须为 bytes=start-end(例如 bytes=0-1024)",
));
}
let api_endpoint = DriveApi::DownloadMedia(self.file_token.clone());
let mut request = api_endpoint
.to_request::<Vec<u8>>()
.query_opt("extra", self.extra);
if let Some(r) = self.range {
request = request.header("Range", &r);
}
let result = Transport::request(request, &self.config, Some(option)).await;
match result {
Ok(response) => {
let data_len = response.data.as_ref().map_or(0, <Vec<u8>>::len);
if data_len > self.max_size {
return Err(openlark_core::error::validation_error(
"max_size",
&format!("下载文件大小 {} 超过限制 {}", data_len, self.max_size),
));
}
Ok(response)
}
Err(e) => Err(e),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_download_media_request() {
let config = Config::default();
let request = DownloadMediaRequest::new(config, "media_token")
.extra("extra")
.range("bytes=0-100");
assert_eq!(request.file_token, "media_token");
assert_eq!(request.extra, Some("extra".to_string()));
assert_eq!(request.range, Some("bytes=0-100".to_string()));
}
#[test]
fn test_valid_range_formats() {
let config = Config::default();
let request1 = DownloadMediaRequest::new(config.clone(), "token").range("bytes=0-1024");
assert_eq!(request1.range, Some("bytes=0-1024".to_string()));
let request2 = DownloadMediaRequest::new(config.clone(), "token").range("bytes=0-0");
assert_eq!(request2.range, Some("bytes=0-0".to_string()));
let request3 = DownloadMediaRequest::new(config.clone(), "token").range("bytes=0-999999");
assert_eq!(request3.range, Some("bytes=0-999999".to_string()));
}
#[test]
fn test_optional_parameters() {
let config = Config::default();
let request1 = DownloadMediaRequest::new(config.clone(), "token");
assert!(request1.extra.is_none());
assert!(request1.range.is_none());
let request2 = DownloadMediaRequest::new(config.clone(), "token").extra("extra_param");
assert_eq!(request2.extra, Some("extra_param".to_string()));
let request3 = DownloadMediaRequest::new(config, "token").range("bytes=0-100");
assert_eq!(request3.range, Some("bytes=0-100".to_string()));
}
#[test]
fn test_download_media_default_max_size() {
let config = Config::default();
let request = DownloadMediaRequest::new(config, "media_token");
assert_eq!(request.max_size, 100 * 1024 * 1024);
}
#[test]
fn test_download_media_custom_max_size() {
let config = Config::default();
let request = DownloadMediaRequest::new(config, "media_token").max_size(512);
assert_eq!(request.max_size, 512);
}
#[tokio::test]
async fn test_download_media_returns_data_on_success() {
use wiremock::MockServer;
use wiremock::matchers::{method, path};
use wiremock::{Mock, ResponseTemplate};
let server = MockServer::start().await;
let body = b"hello download binary payload".to_vec();
Mock::given(method("GET"))
.and(path("/open-apis/drive/v1/medias/media_token_001/download"))
.respond_with(ResponseTemplate::new(200).set_body_bytes(body.clone()))
.mount(&server)
.await;
let config = Config::builder()
.app_id("ci_app_id")
.app_secret("ci_app_secret")
.base_url(server.uri())
.enable_token_cache(false)
.build();
let resp = DownloadMediaRequest::new(config, "media_token_001")
.execute()
.await
.expect("下载素材应成功");
let data = resp.data.expect("响应应包含二进制数据");
assert_eq!(data, body);
let received = server.received_requests().await.unwrap_or_default();
assert_eq!(received.len(), 1);
assert_eq!(
received[0].url.path(),
"/open-apis/drive/v1/medias/media_token_001/download"
);
}
}