Genarrative/server-rs/crates/platform-image/tests/vector_engine.rs

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 = "<html><head><title>502 Bad Gateway</title></head><body><center><h1>502 Bad Gateway</h1></center><hr><center>nginx</center></body></html>";
                    let response = format!(
                        "HTTP/1.1 502 Bad Gateway\r\nContent-Type: text/html\r\nContent-Length: {}\r\n\r\n{}",
                        body.len(),
                        body
                    );
                    let _ = stream.write_all(response.as_bytes()).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: 1_000,
    };
    let http_client =
        build_vector_engine_image_http_client(&settings).expect("client should build");

    let generated = create_vector_engine_image_generation(
        &http_client,
        &settings,
        "测试提示词",
        None,
        "1024x1024",
        1,
        &[],
        "测试 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();
}