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, build_vector_engine_image_request_body_with_model, build_vector_engine_nanobanana_generate_content_request_body, create_vector_engine_image_edit, create_vector_engine_image_generation, create_vector_engine_nanobanana_generate_content, vector_engine_images_edit_url, vector_engine_images_generation_url, vector_engine_nanobanana_generate_content_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" ); } #[test] fn vector_engine_normalizes_2k_landscape_spec_size() { let body = build_vector_engine_image_request_body("生成规范图", None, "2048x1152", 1, &[]); assert_eq!(body["model"], GPT_IMAGE_2_MODEL); assert_eq!(body["size"], "2048x1152"); assert_eq!(body["n"], 1); } #[test] fn vector_engine_request_body_can_use_nanobanana2_model() { let body = build_vector_engine_image_request_body_with_model( "gemini-3.1-flash-image-preview", "生成图标 spritesheet", None, "512x512", 1, &[], ); assert_eq!(body["model"], "gemini-3.1-flash-image-preview"); assert_eq!(body["size"], "512x512"); assert_eq!(body["n"], 1); } #[test] fn vector_engine_request_body_can_use_nanobanana2_half_k() { let body = build_vector_engine_image_request_body_with_model( "gemini-3.1-flash-image-preview", "生成图标 spritesheet", None, "512", 1, &[], ); assert_eq!(body["model"], "gemini-3.1-flash-image-preview"); assert_eq!(body["size"], "512"); } #[test] fn nanobanana_generate_content_body_carries_aspect_ratio_and_image_size() { let body = build_vector_engine_nanobanana_generate_content_request_body( "生成角色图", Some("文字、水印"), "2:3", "512", &[], ); assert_eq!(body["contents"][0]["role"], "user"); assert_eq!( body["contents"][0]["parts"][0]["text"], "生成角色图\n避免:文字、水印" ); assert_eq!(body["generationConfig"]["responseModalities"][0], "IMAGE"); assert_eq!( body["generationConfig"]["imageConfig"]["aspectRatio"], "2:3" ); assert_eq!(body["generationConfig"]["imageConfig"]["imageSize"], "512"); assert!(body.get("model").is_none()); assert!(body.get("n").is_none()); } #[test] fn nanobanana_generate_content_url_uses_model_path() { let settings = VectorEngineImageSettings { base_url: "https://vector.example/v1".to_string(), api_key: "test-key".to_string(), request_timeout_ms: 1_000, }; assert_eq!( vector_engine_nanobanana_generate_content_url(&settings, "gemini-3.1-flash-image-preview"), "https://vector.example/v1beta/models/gemini-3.1-flash-image-preview:generateContent" ); } #[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(); } #[tokio::test] async fn nanobanana_generate_content_posts_native_body_and_reads_inline_data() { 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 server = tokio::spawn(async move { let Ok((mut stream, _)) = listener.accept().await else { return; }; let mut request = Vec::new(); let mut buffer = [0_u8; 4096]; loop { let Ok(read) = stream.read(&mut buffer).await else { return; }; if read == 0 { return; } request.extend_from_slice(&buffer[..read]); if request.windows(4).any(|window| window == b"\r\n\r\n") { break; } } let request_text = String::from_utf8_lossy(request.as_slice()); assert!( request_text.contains("/v1beta/models/gemini-3.1-flash-image-preview:generateContent") ); assert!(request_text.contains("\"aspectRatio\":\"2:3\"")); assert!(request_text.contains("\"imageSize\":\"512\"")); let body = r#"{"candidates":[{"content":{"parts":[{"inlineData":{"mimeType":"image/png","data":"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), api_key: "test-key".to_string(), request_timeout_ms: 1_000, }; let client = build_vector_engine_image_http_client(&settings).expect("client should build"); let generated = create_vector_engine_nanobanana_generate_content( &client, &settings, "gemini-3.1-flash-image-preview", "生成角色图", Some("文字、水印"), "2:3", "512", &[], "测试 nanobanana", ) .await .expect("nanobanana response should parse"); assert_eq!(generated.images.len(), 1); assert_eq!(generated.images[0].mime_type, "image/png"); server.abort(); } #[tokio::test] async fn vector_engine_image_generation_retries_upstream_502_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 { let body = "