112 lines
3.7 KiB
Rust
112 lines
3.7 KiB
Rust
use platform_image::vector_engine::{
|
|
GPT_IMAGE_2_MODEL, ReferenceImage, VECTOR_ENGINE_PROVIDER, VectorEngineImageSettings,
|
|
build_vector_engine_image_http_client, build_vector_engine_image_request_body,
|
|
create_vector_engine_image_edit, vector_engine_images_edit_url,
|
|
vector_engine_images_generation_url,
|
|
};
|
|
use std::{
|
|
sync::{
|
|
Arc,
|
|
atomic::{AtomicUsize, Ordering},
|
|
},
|
|
time::Duration,
|
|
};
|
|
use tokio::{
|
|
io::{AsyncReadExt, AsyncWriteExt},
|
|
net::TcpListener,
|
|
};
|
|
|
|
#[test]
|
|
fn vector_engine_module_exposes_provider_protocol_helpers() {
|
|
let settings = VectorEngineImageSettings {
|
|
base_url: "https://vector.example/v1".to_string(),
|
|
api_key: "test-key".to_string(),
|
|
request_timeout_ms: 1_000,
|
|
};
|
|
|
|
let body =
|
|
build_vector_engine_image_request_body("雾海神殿", Some("文字，水印"), "16:9", 9, &[]);
|
|
|
|
assert_eq!(GPT_IMAGE_2_MODEL, "gpt-image-2");
|
|
assert_eq!(VECTOR_ENGINE_PROVIDER, "vector-engine");
|
|
assert_eq!(body["model"], GPT_IMAGE_2_MODEL);
|
|
assert_eq!(body["size"], "1536x1024");
|
|
assert_eq!(body["n"], 4);
|
|
assert_eq!(body["prompt"], "雾海神殿\n避免：文字，水印");
|
|
assert_eq!(
|
|
vector_engine_images_generation_url(&settings),
|
|
"https://vector.example/v1/images/generations"
|
|
);
|
|
assert_eq!(
|
|
vector_engine_images_edit_url(&settings),
|
|
"https://vector.example/v1/images/edits"
|
|
);
|
|
}
|
|
|
|
#[tokio::test]
|
|
async fn vector_engine_image_edit_retries_send_timeout_once_and_succeeds() {
|
|
let listener = TcpListener::bind("127.0.0.1:0")
|
|
.await
|
|
.expect("mock server should bind");
|
|
let server_addr = listener
|
|
.local_addr()
|
|
.expect("mock server address should be readable");
|
|
let request_count = Arc::new(AtomicUsize::new(0));
|
|
let request_count_for_server = Arc::clone(&request_count);
|
|
|
|
let server = tokio::spawn(async move {
|
|
loop {
|
|
let Ok((mut stream, _)) = listener.accept().await else {
|
|
break;
|
|
};
|
|
let request_index = request_count_for_server.fetch_add(1, Ordering::SeqCst);
|
|
tokio::spawn(async move {
|
|
let mut buffer = [0_u8; 4096];
|
|
let _ = stream.read(&mut buffer).await;
|
|
if request_index == 0 {
|
|
tokio::time::sleep(Duration::from_millis(120)).await;
|
|
return;
|
|
}
|
|
|
|
let body = r#"{"data":[{"b64_json":"iVBORw0KGgpyZXN0"}]}"#;
|
|
let response = format!(
|
|
"HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: {}\r\n\r\n{}",
|
|
body.len(),
|
|
body
|
|
);
|
|
let _ = stream.write_all(response.as_bytes()).await;
|
|
});
|
|
}
|
|
});
|
|
|
|
let settings = VectorEngineImageSettings {
|
|
base_url: format!("http://{server_addr}/v1"),
|
|
api_key: "test-key".to_string(),
|
|
request_timeout_ms: 40,
|
|
};
|
|
let http_client =
|
|
build_vector_engine_image_http_client(&settings).expect("client should build");
|
|
let reference_image = ReferenceImage {
|
|
bytes: b"reference".to_vec(),
|
|
mime_type: "image/png".to_string(),
|
|
file_name: "reference.png".to_string(),
|
|
};
|
|
|
|
let generated = create_vector_engine_image_edit(
|
|
&http_client,
|
|
&settings,
|
|
"测试提示词",
|
|
None,
|
|
"1024x1024",
|
|
&reference_image,
|
|
"测试 VectorEngine 图片编辑失败",
|
|
)
|
|
.await
|
|
.expect("second attempt should return generated image");
|
|
|
|
assert_eq!(generated.images.len(), 1);
|
|
assert_eq!(generated.images[0].mime_type, "image/png");
|
|
assert_eq!(request_count.load(Ordering::SeqCst), 2);
|
|
server.abort();
|
|
}
|