Genarrative/server-rs/crates/platform-audio/tests/elevenlabs.rs

use std::{
    io::{Read, Write},
    net::{TcpListener, TcpStream},
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
        mpsc,
    },
    thread,
    time::{Duration, Instant},
};

use platform_audio::{
    AudioError, ELEVENLABS_PROVIDER, ELEVENLABS_SOUND_EFFECT_MODEL, ElevenLabsAudioSettings,
    ElevenLabsSoundEffectRequest, MAX_GENERATED_AUDIO_BYTES, build_elevenlabs_audio_http_client,
    generate_elevenlabs_sound_effect,
};

const TEST_API_KEY: &str = "elevenlabs-test-key";
const TEST_MP3: &[u8] = include_bytes!("../../../../public/wooden-fish/default-hit-sound.mp3");

fn runtime() -> tokio::runtime::Runtime {
    tokio::runtime::Builder::new_current_thread()
        .enable_all()
        .build()
        .expect("tokio runtime should build")
}

fn settings(base_url: String, request_timeout_ms: u64) -> ElevenLabsAudioSettings {
    ElevenLabsAudioSettings {
        base_url,
        api_key: TEST_API_KEY.to_string(),
        request_timeout_ms,
        request_deadline: None,
    }
}

fn request() -> ElevenLabsSoundEffectRequest {
    ElevenLabsSoundEffectRequest {
        text: "crisp metallic coin pickup".to_string(),
        duration_seconds: Some(1.234_567_89),
        loop_enabled: true,
    }
}

fn read_http_request(stream: &mut TcpStream) -> Vec<u8> {
    stream
        .set_read_timeout(Some(Duration::from_secs(2)))
        .expect("mock request read timeout should configure");
    let mut request_bytes = Vec::new();
    let mut buffer = [0_u8; 4096];
    let mut expected_len = None;
    loop {
        let read = stream.read(&mut buffer).expect("mock request should read");
        if read == 0 {
            break;
        }
        request_bytes.extend_from_slice(&buffer[..read]);
        if expected_len.is_none()
            && let Some(header_end) = find_header_end(&request_bytes)
        {
            let headers = String::from_utf8_lossy(&request_bytes[..header_end]);
            let content_length = headers
                .lines()
                .find_map(|line| {
                    line.split_once(':').and_then(|(name, value)| {
                        name.eq_ignore_ascii_case("content-length")
                            .then(|| value.trim().parse::<usize>().ok())
                            .flatten()
                    })
                })
                .unwrap_or(0);
            expected_len = Some(header_end + 4 + content_length);
        }
        if expected_len.is_some_and(|length| request_bytes.len() >= length) {
            break;
        }
    }
    request_bytes
}

fn find_header_end(bytes: &[u8]) -> Option<usize> {
    bytes.windows(4).position(|window| window == b"\r\n\r\n")
}

fn spawn_single_response_server(
    response: Vec<u8>,
) -> (String, mpsc::Receiver<Vec<u8>>, thread::JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("mock server should bind");
    let address = listener
        .local_addr()
        .expect("mock address should be readable");
    let (request_sender, request_receiver) = mpsc::channel();
    let server = thread::spawn(move || {
        let (mut stream, _) = listener.accept().expect("mock request should connect");
        let request_bytes = read_http_request(&mut stream);
        request_sender
            .send(request_bytes)
            .expect("captured request should send");
        stream
            .write_all(&response)
            .expect("mock response should write");
    });
    (format!("http://{address}"), request_receiver, server)
}

fn spawn_counting_response_server(
    response: Vec<u8>,
    response_delay: Duration,
) -> (String, Arc<AtomicUsize>, thread::JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("mock server should bind");
    let address = listener
        .local_addr()
        .expect("mock address should be readable");
    let request_count = Arc::new(AtomicUsize::new(0));
    let request_count_for_server = Arc::clone(&request_count);
    let server = thread::spawn(move || {
        let (mut stream, _) = listener.accept().expect("mock request should connect");
        request_count_for_server.fetch_add(1, Ordering::SeqCst);
        let _ = read_http_request(&mut stream);
        thread::sleep(response_delay);
        let _ = stream.write_all(&response);
        drop(stream);

        listener
            .set_nonblocking(true)
            .expect("mock listener should become nonblocking");
        let deadline = Instant::now() + Duration::from_millis(150);
        while Instant::now() < deadline {
            match listener.accept() {
                Ok((mut retry_stream, _)) => {
                    request_count_for_server.fetch_add(1, Ordering::SeqCst);
                    let _ = read_http_request(&mut retry_stream);
                    let _ = retry_stream.write_all(&response);
                }
                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                    thread::sleep(Duration::from_millis(5));
                }
                Err(error) => panic!("mock listener failed: {error}"),
            }
        }
    });
    (format!("http://{address}"), request_count, server)
}

/// 用于断言「provider 一次都没被调用」：只做非阻塞轮询，不会在零请求时把 join 卡死。
fn spawn_never_answering_server() -> (String, Arc<AtomicUsize>, thread::JoinHandle<()>) {
    let listener = TcpListener::bind("127.0.0.1:0").expect("mock server should bind");
    let address = listener
        .local_addr()
        .expect("mock address should be readable");
    let request_count = Arc::new(AtomicUsize::new(0));
    let request_count_for_server = Arc::clone(&request_count);
    let server = thread::spawn(move || {
        listener
            .set_nonblocking(true)
            .expect("mock listener should become nonblocking");
        let deadline = Instant::now() + Duration::from_millis(200);
        while Instant::now() < deadline {
            match listener.accept() {
                Ok((mut stream, _)) => {
                    request_count_for_server.fetch_add(1, Ordering::SeqCst);
                    let _ = read_http_request(&mut stream);
                }
                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                    thread::sleep(Duration::from_millis(5));
                }
                Err(error) => panic!("mock listener failed: {error}"),
            }
        }
    });
    (format!("http://{address}"), request_count, server)
}

fn success_response(content_type: Option<&str>, body: &[u8]) -> Vec<u8> {
    let content_type_header = content_type
        .map(|value| format!("Content-Type: {value}\r\n"))
        .unwrap_or_default();
    let mut response = format!(
        "HTTP/1.1 200 OK\r\n{content_type_header}Content-Length: {}\r\nConnection: close\r\n\r\n",
        body.len()
    )
    .into_bytes();
    response.extend_from_slice(body);
    response
}

#[test]
fn direct_client_sends_the_fixed_contract_and_returns_probed_mp3() {
    let (base_url, captured_request, server) =
        spawn_single_response_server(success_response(Some("audio/mpeg"), TEST_MP3));
    let settings = settings(format!("{base_url}/v1/"), 2_000);
    let client =
        build_elevenlabs_audio_http_client(&settings).expect("ElevenLabs HTTP client should build");

    let generated = runtime()
        .block_on(generate_elevenlabs_sound_effect(
            &client,
            &settings,
            request(),
        ))
        .expect("valid MP3 response should pass");
    let captured_request = captured_request
        .recv_timeout(Duration::from_secs(2))
        .expect("mock request should be captured");
    server.join().expect("mock server should finish");

    assert_eq!(generated.audio.bytes, TEST_MP3);
    assert_eq!(generated.audio.mime_type, "audio/mpeg");
    assert_eq!(generated.audio.extension, "mp3");
    assert!((0.65..=0.75).contains(&generated.duration_seconds));

    let header_end = find_header_end(&captured_request).expect("request headers should terminate");
    let headers = String::from_utf8_lossy(&captured_request[..header_end]).to_ascii_lowercase();
    let body: serde_json::Value = serde_json::from_slice(&captured_request[header_end + 4..])
        .expect("request body should be JSON");
    assert!(
        headers.starts_with("post /v1/sound-generation?output_format=mp3_44100_128 http/1.1\r\n")
    );
    assert!(headers.contains("xi-api-key: elevenlabs-test-key\r\n"));
    assert!(headers.contains("accept-encoding: identity\r\n"));
    assert_eq!(body["text"], "crisp metallic coin pickup");
    assert_eq!(body["model_id"], ELEVENLABS_SOUND_EFFECT_MODEL);
    assert_eq!(body["duration_seconds"], 1.234_567_89);
    assert_eq!(body["loop"], true);
    assert_eq!(body["prompt_influence"], 0.3);
    assert!(!String::from_utf8_lossy(&captured_request[header_end + 4..]).contains(TEST_API_KEY));
}

#[test]
fn octet_stream_and_missing_content_type_require_and_accept_real_mp3() {
    for content_type in [Some("audio/mp3"), Some("application/octet-stream"), None] {
        let (base_url, _captured_request, server) =
            spawn_single_response_server(success_response(content_type, TEST_MP3));
        let settings = settings(base_url, 2_000);
        let client = build_elevenlabs_audio_http_client(&settings)
            .expect("ElevenLabs HTTP client should build");

        let generated = runtime()
            .block_on(generate_elevenlabs_sound_effect(
                &client,
                &settings,
                request(),
            ))
            .expect("real MP3 should pass MIME fallback");
        assert_eq!(generated.audio.mime_type, "audio/mpeg");
        server.join().expect("mock server should finish");
    }
}

#[test]
fn explicit_wrong_mime_empty_and_corrupt_bodies_are_rejected() {
    let cases = [
        (
            success_response(Some("text/html"), TEST_MP3),
            "invalid_audio",
        ),
        (
            success_response(Some("application/json"), br#"{"error":"no"}"#),
            "invalid_audio",
        ),
        (success_response(Some("audio/mpeg"), b""), "invalid_audio"),
        (
            success_response(Some("audio/mpeg"), b"ID3-corrupt"),
            "duration_probe_failed",
        ),
    ];
    for (response, expected_reason) in cases {
        let (base_url, _captured_request, server) = spawn_single_response_server(response);
        let settings = settings(base_url, 2_000);
        let client = build_elevenlabs_audio_http_client(&settings)
            .expect("ElevenLabs HTTP client should build");

        let error = runtime()
            .block_on(generate_elevenlabs_sound_effect(
                &client,
                &settings,
                request(),
            ))
            .expect_err("invalid audio response should fail");
        assert_eq!(error.provider(), ELEVENLABS_PROVIDER);
        assert_eq!(error.reason_code(), Some(expected_reason));
        server.join().expect("mock server should finish");
    }
}

#[test]
fn redirects_http_failures_and_body_failures_do_not_retry_the_provider_post() {
    let responses = [
        b"HTTP/1.1 307 Temporary Redirect\r\nLocation: /v1/sound-generation\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_vec(),
        b"HTTP/1.1 401 Unauthorized\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_vec(),
        b"HTTP/1.1 422 Unprocessable Entity\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_vec(),
        b"HTTP/1.1 429 Too Many Requests\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_vec(),
        b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 29\r\nConnection: close\r\n\r\nprovider-secret-response-body".to_vec(),
        b"HTTP/1.1 200 OK\r\nContent-Type: audio/mpeg\r\nContent-Length: 20\r\nConnection: close\r\n\r\nID3".to_vec(),
    ];
    for response in responses {
        let (base_url, request_count, server) =
            spawn_counting_response_server(response, Duration::ZERO);
        let settings = settings(base_url, 500);
        let client = build_elevenlabs_audio_http_client(&settings)
            .expect("ElevenLabs HTTP client should build");

        let error = runtime()
            .block_on(generate_elevenlabs_sound_effect(
                &client,
                &settings,
                request(),
            ))
            .expect_err("provider failure should fail once");
        server.join().expect("mock server should finish");

        assert_eq!(error.provider(), ELEVENLABS_PROVIDER);
        assert_eq!(request_count.load(Ordering::SeqCst), 1);
        assert!(!format!("{error:?}").contains(TEST_API_KEY));
        assert!(!format!("{error:?}").contains("provider-secret-response-body"));
    }
}

#[test]
fn timeout_does_not_retry_the_provider_post() {
    let response = success_response(Some("audio/mpeg"), TEST_MP3);
    let (base_url, request_count, server) =
        spawn_counting_response_server(response, Duration::from_millis(100));
    let settings = settings(base_url, 20);
    let client =
        build_elevenlabs_audio_http_client(&settings).expect("ElevenLabs HTTP client should build");

    let error = runtime()
        .block_on(generate_elevenlabs_sound_effect(
            &client,
            &settings,
            request(),
        ))
        .expect_err("provider timeout should fail once");
    server.join().expect("mock server should finish");

    assert!(matches!(error, AudioError::Request { timeout: true, .. }));
    assert_eq!(request_count.load(Ordering::SeqCst), 1);
}

#[test]
fn an_exhausted_provider_budget_fails_before_any_provider_post() {
    let (base_url, request_count, server) = spawn_never_answering_server();
    let mut settings = settings(base_url, 180_000);
    settings.request_deadline = Some(Instant::now() - Duration::from_millis(1));
    let client =
        build_elevenlabs_audio_http_client(&settings).expect("ElevenLabs HTTP client should build");

    let error = runtime()
        .block_on(generate_elevenlabs_sound_effect(
            &client,
            &settings,
            request(),
        ))
        .expect_err("an exhausted provider budget should fail before sending");
    server.join().expect("mock server should finish");

    // 预算耗尽时 provider 一次都不能被调用，否则慢翻译会白白消耗一次计费请求。
    assert_eq!(request_count.load(Ordering::SeqCst), 0);
    assert_eq!(error.provider(), ELEVENLABS_PROVIDER);
    assert!(matches!(error, AudioError::Request { timeout: true, .. }));
}

#[test]
fn a_remaining_provider_budget_still_sends_exactly_one_provider_post() {
    let response = success_response(Some("audio/mpeg"), TEST_MP3);
    let (base_url, request_count, server) =
        spawn_counting_response_server(response, Duration::ZERO);
    let mut settings = settings(base_url, 180_000);
    settings.request_deadline = Some(Instant::now() + Duration::from_secs(30));
    let client =
        build_elevenlabs_audio_http_client(&settings).expect("ElevenLabs HTTP client should build");

    let generated = runtime()
        .block_on(generate_elevenlabs_sound_effect(
            &client,
            &settings,
            request(),
        ))
        .expect("a remaining provider budget should still generate");
    server.join().expect("mock server should finish");

    assert_eq!(request_count.load(Ordering::SeqCst), 1);
    assert!(generated.duration_seconds > 0.0);
}

#[test]
fn content_length_precheck_accepts_the_limit_and_rejects_limit_plus_one() {
    for (content_length, should_be_size_error) in [
        (MAX_GENERATED_AUDIO_BYTES, false),
        (MAX_GENERATED_AUDIO_BYTES + 1, true),
    ] {
        let response = format!(
            "HTTP/1.1 200 OK\r\nContent-Type: audio/mpeg\r\nContent-Length: {content_length}\r\nConnection: close\r\n\r\nID3"
        )
        .into_bytes();
        let (base_url, request_count, server) =
            spawn_counting_response_server(response, Duration::ZERO);
        let settings = settings(base_url, 500);
        let client = build_elevenlabs_audio_http_client(&settings)
            .expect("ElevenLabs HTTP client should build");

        let error = runtime()
            .block_on(generate_elevenlabs_sound_effect(
                &client,
                &settings,
                request(),
            ))
            .expect_err("incomplete or oversized response should fail");
        server.join().expect("mock server should finish");

        assert_eq!(
            error.message().contains("超过 40 MiB"),
            should_be_size_error
        );
        if should_be_size_error {
            assert_eq!(error.reason_code(), Some("invalid_audio"));
        }
        assert_eq!(request_count.load(Ordering::SeqCst), 1);
    }
}

#[test]
fn chunked_limit_plus_one_is_rejected_without_unbounded_allocation() {
    let listener = TcpListener::bind("127.0.0.1:0").expect("mock server should bind");
    let address = listener
        .local_addr()
        .expect("mock address should be readable");
    let server = thread::spawn(move || {
        let (mut stream, _) = listener.accept().expect("mock request should connect");
        let _ = read_http_request(&mut stream);
        stream
            .write_all(
                b"HTTP/1.1 200 OK\r\nContent-Type: audio/mpeg\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n",
            )
            .expect("chunked response headers should write");
        let chunk = vec![0_u8; 64 * 1024];
        for _ in 0..(MAX_GENERATED_AUDIO_BYTES / chunk.len()) {
            if stream.write_all(b"10000\r\n").is_err()
                || stream.write_all(&chunk).is_err()
                || stream.write_all(b"\r\n").is_err()
            {
                return;
            }
        }
        let _ = stream.write_all(b"1\r\nX\r\n0\r\n\r\n");
    });
    let settings = settings(format!("http://{address}"), 5_000);
    let client =
        build_elevenlabs_audio_http_client(&settings).expect("ElevenLabs HTTP client should build");

    let error = runtime()
        .block_on(generate_elevenlabs_sound_effect(
            &client,
            &settings,
            request(),
        ))
        .expect_err("chunked limit plus one should fail");
    server.join().expect("mock server should finish");

    assert!(error.message().contains("超过 40 MiB"));
    assert_eq!(error.reason_code(), Some("invalid_audio"));
}