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

use platform_image::vector_engine::{
    GPT_IMAGE_2_MODEL, PlatformImageError, 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, Instant},
};
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,
        request_deadline: None,
    };

    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_clamps_gpt_image_2_explicit_pixel_sizes_to_its_supported_pixel_budget() {
    let cover = build_vector_engine_image_request_body("宣发首图", None, "720x540", 1, &[]);
    let detail = build_vector_engine_image_request_body("详情单图", None, "720x1280", 1, &[]);
    let poster = build_vector_engine_image_request_body("运营海报", None, "1280x720", 1, &[]);

    assert_eq!(cover["size"], "944x704");
    assert_eq!(detail["size"], "720x1280");
    assert_eq!(poster["size"], "1280x720");
}

#[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_only_enforces_the_gpt_image_2_pixel_budget_for_that_model() {
    let gpt_body = build_vector_engine_image_request_body_with_model(
        GPT_IMAGE_2_MODEL,
        "小尺寸图",
        None,
        "640x640",
        1,
        &[],
    );
    let nanobanana_body = build_vector_engine_image_request_body_with_model(
        "gemini-3.1-flash-image-preview",
        "小尺寸图",
        None,
        "640x640",
        1,
        &[],
    );
    let oversized_gpt_body = build_vector_engine_image_request_body_with_model(
        GPT_IMAGE_2_MODEL,
        "大尺寸图",
        None,
        "4096x4096",
        1,
        &[],
    );

    assert_eq!(gpt_body["size"], "816x816");
    assert_eq!(nanobanana_body["size"], "640x640");
    assert_eq!(oversized_gpt_body["size"], "2880x2880");
}

#[test]
fn vector_engine_gpt_image_2_sizes_always_meet_the_full_provider_envelope() {
    for size in [
        "1x1",
        "720x540",
        "3841x1280",
        "4096x4096",
        "3200x400",
        "16x4096",
        "3840x3840",
    ] {
        let body = build_vector_engine_image_request_body("约束测试", None, size, 1, &[]);
        let normalized = body["size"].as_str().expect("size should be a string");
        let (width, height) = normalized
            .split_once('x')
            .expect("gpt-image-2 size should be explicit pixels");
        let width = width.parse::<u32>().expect("width should be numeric");
        let height = height.parse::<u32>().expect("height should be numeric");
        let pixels = u64::from(width) * u64::from(height);

        assert!(width <= 3_840 && height <= 3_840, "{size} -> {normalized}");
        assert!(width.is_multiple_of(16) && height.is_multiple_of(16));
        assert!((655_360..=8_294_400).contains(&pixels));
        assert!(width.max(height) <= width.min(height) * 3);
    }
}

#[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,
        request_deadline: None,
    };

    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,
        request_deadline: None,
    };
    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 vector_engine_deadline_clips_stalled_attempt_and_prevents_retry() {
    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;
            };
            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;
                tokio::time::sleep(Duration::from_secs(1)).await;
            });
        }
    });

    let started_at = Instant::now();
    let settings = VectorEngineImageSettings {
        base_url: format!("http://{server_addr}/v1"),
        api_key: "test-key".to_string(),
        request_timeout_ms: 5_000,
        request_deadline: Some(started_at + Duration::from_millis(150)),
    };
    let http_client =
        build_vector_engine_image_http_client(&settings).expect("client should build");

    let error = create_vector_engine_image_generation(
        &http_client,
        &settings,
        "测试提示词",
        None,
        "1024x1024",
        1,
        &[],
        "测试 VectorEngine 图片生成失败",
    )
    .await
    .expect_err("stalled request should exhaust the shared deadline");

    assert!(matches!(
        error,
        PlatformImageError::Request { timeout: true, .. }
    ));
    assert!(
        started_at.elapsed() < Duration::from_secs(1),
        "attempt 应使用剩余 deadline，而不是完整配置 timeout"
    );
    assert_eq!(request_count.load(Ordering::SeqCst), 1);
    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,
        request_deadline: None,
    };
    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,
        request_deadline: None,
    };
    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();
}