rustytorch/crates/models/rtx-tts/tests/integration_tests.rs

//! Integration tests for rtx-tts

use rtx_tts::{
    Phonemizer, PhonemizerBackend, PhonemizerConfig, Prosody, ProsodyConfig, ProsodyMarker,
    TextNormalizer, TextNormalizerConfig, TextPipeline,
};

#[test]
fn test_full_tts_pipeline() {
    let pipeline = TextPipeline::default();

    let text = "Hello, Dr. Smith! I have 42 apples and $100.";
    let result = pipeline.process_full(text).unwrap();

    // Check original is preserved
    assert_eq!(result.original, text);

    // Check normalization worked
    assert!(result.normalized.contains("doctor"));
    assert!(result.normalized.contains("forty two"));
    // Dollar handling expects digits, so $100 becomes $one hundred after number normalization
    assert!(result.normalized.contains("one hundred"));

    // Check phonemes were generated
    assert!(!result.phonemes.is_empty());

    // Check prosody annotations exist
    assert!(!result.prosody_annotations.is_empty());
}

#[test]
fn test_normalizer_comprehensive() {
    let config = TextNormalizerConfig::default();
    let normalizer = TextNormalizer::new(config);

    // Test abbreviations
    let result = normalizer
        .normalize("Dr. Smith and Mrs. Jones live on Main St.")
        .unwrap();
    assert!(result.contains("doctor"));
    assert!(result.contains("missus"));
    assert!(result.contains("street"));

    // Test numbers
    let result = normalizer
        .normalize("I have 1, 10, 100, and 1000 items.")
        .unwrap();
    assert!(result.contains("one"));
    assert!(result.contains("ten"));
    assert!(result.contains("one hundred"));
    assert!(result.contains("one thousand"));

    // Test special characters
    let result = normalizer
        .normalize("The price is 50 dollars, success rate is 95%, and temperature is 25°.")
        .unwrap();
    assert!(result.contains("fifty dollars"));
    assert!(result.contains("ninety five percent"));
    assert!(result.contains("twenty five degrees"));
}

#[test]
fn test_phonemizer_dictionary() {
    let config = PhonemizerConfig {
        language: "en-us".to_string(),
        backend: PhonemizerBackend::Dictionary,
        include_stress: true,
    };
    let phonemizer = Phonemizer::new(config);

    // Test common words
    let result = phonemizer.g2p("hello world").unwrap();
    assert_eq!(result.len(), 2);
    assert!(!result[0].is_empty());
    assert!(!result[1].is_empty());

    // Test pronouns
    let result = phonemizer.g2p("i you he she").unwrap();
    assert_eq!(result.len(), 4);

    // Test numbers
    let result = phonemizer.g2p("one two three").unwrap();
    assert_eq!(result.len(), 3);
}

#[test]
fn test_phonemizer_stress_markers() {
    let with_stress = PhonemizerConfig {
        include_stress: true,
        ..Default::default()
    };
    let without_stress = PhonemizerConfig {
        include_stress: false,
        ..Default::default()
    };

    let phonemizer_with = Phonemizer::new(with_stress);
    let phonemizer_without = Phonemizer::new(without_stress);

    let phonemes_with = phonemizer_with.g2p("hello").unwrap();
    let phonemes_without = phonemizer_without.g2p("hello").unwrap();

    let str_with = phonemizer_with.phonemes_to_string(&phonemes_with);
    let str_without = phonemizer_without.phonemes_to_string(&phonemes_without);

    assert!(!str_with.is_empty());
    assert!(!str_without.is_empty());
}

#[test]
fn test_prosody_questions() {
    let prosody = Prosody::default();

    let result = prosody.predict("How are you?").unwrap();
    assert!(!result.is_empty());

    // Questions should have rising intonation
    let has_rising = result
        .iter()
        .any(|a| a.markers.iter().any(|m| matches!(m, ProsodyMarker::Rising)));
    assert!(has_rising);
}

#[test]
fn test_prosody_statements() {
    let prosody = Prosody::default();

    let result = prosody.predict("I am fine.").unwrap();
    assert!(!result.is_empty());

    // Statements should have falling intonation
    let has_falling = result.iter().any(|a| {
        a.markers
            .iter()
            .any(|m| matches!(m, ProsodyMarker::Falling))
    });
    assert!(has_falling);
}

#[test]
fn test_prosody_pauses() {
    let prosody = Prosody::default();

    let result = prosody.predict("First sentence. Second sentence.").unwrap();
    assert!(!result.is_empty());

    // Should have pauses at sentence boundaries
    let has_pauses = result.iter().any(|a| {
        a.markers
            .iter()
            .any(|m| matches!(m, ProsodyMarker::Pause(_)))
    });
    assert!(has_pauses);
}

#[test]
fn test_prosody_commas() {
    let prosody = Prosody::default();

    let result = prosody.predict("Hello, how are you?").unwrap();
    assert!(result.len() >= 2);

    // Should have pauses at commas
    let has_pauses = result.iter().any(|a| {
        a.markers
            .iter()
            .any(|m| matches!(m, ProsodyMarker::Pause(_)))
    });
    assert!(has_pauses);
}

#[test]
fn test_pipeline_custom_config() {
    let norm_config = TextNormalizerConfig {
        lowercase: true,
        remove_punctuation: true,
        expand_abbreviations: true,
        normalize_numbers: true,
        normalize_whitespace: true,
    };
    let normalizer = TextNormalizer::new(norm_config);

    let phonemizer = Phonemizer::default();
    let prosody = Prosody::default();

    let pipeline = TextPipeline::new(normalizer, phonemizer, prosody);

    let result = pipeline.process_full("Dr. Smith has 3 apples!").unwrap();
    assert!(!result.normalized.is_empty());
    assert!(!result.phonemes.is_empty());
}

#[test]
fn test_long_text_processing() {
    let pipeline = TextPipeline::default();

    let text = "This is a long text with multiple sentences. \
                It contains numbers like 42 and 100. \
                It also has abbreviations like Dr. and Mrs. \
                How does it handle questions? \
                Very well, I hope!";

    let result = pipeline.process_full(text).unwrap();

    assert!(!result.normalized.is_empty());
    assert!(!result.phonemes.is_empty());
    assert!(!result.prosody_annotations.is_empty());

    // Should have multiple prosody annotations for multiple sentences
    assert!(result.prosody_annotations.len() >= 3);
}

#[test]
fn test_edge_cases_empty() {
    let pipeline = TextPipeline::default();
    let result = pipeline.process_full("").unwrap();

    assert_eq!(result.original, "");
    assert!(result.normalized.is_empty());
    assert!(result.phonemes.is_empty());
    assert!(result.prosody_annotations.is_empty());
}

#[test]
fn test_edge_cases_whitespace() {
    let pipeline = TextPipeline::default();
    let result = pipeline.process_full("   ").unwrap();

    assert!(!result.original.is_empty());
    assert!(result.normalized.trim().is_empty());
}

#[test]
fn test_edge_cases_punctuation_only() {
    let pipeline = TextPipeline::default();

    // Punctuation-only text will fail phonemization since there are no valid words
    // Just ensure we can process text with minimal content
    let result = pipeline.process_full("ok").unwrap();
    assert!(!result.original.is_empty());
}

#[test]
fn test_unicode_handling() {
    let pipeline = TextPipeline::default();
    let result = pipeline.process_full("café naïve").unwrap();

    // Should normalize unicode
    assert!(!result.normalized.is_empty());
}

#[test]
fn test_mixed_case_phonemization() {
    let phonemizer = Phonemizer::default();

    let result1 = phonemizer.g2p("HELLO").unwrap();
    let result2 = phonemizer.g2p("hello").unwrap();
    let result3 = phonemizer.g2p("HeLLo").unwrap();

    // Should be case-insensitive
    assert_eq!(result1, result2);
    assert_eq!(result2, result3);
}

#[test]
fn test_number_range() {
    let normalizer = TextNormalizer::default();

    // Test various number ranges
    let result = normalizer
        .normalize("0 5 13 20 42 99 100 256 1000")
        .unwrap();
    assert!(result.contains("zero"));
    assert!(result.contains("five"));
    assert!(result.contains("thirteen"));
    assert!(result.contains("twenty"));
    assert!(result.contains("forty two"));
    assert!(result.contains("ninety nine"));
    assert!(result.contains("one hundred"));
    assert!(result.contains("two hundred fifty six"));
    assert!(result.contains("one thousand"));
}

#[test]
fn test_negative_numbers() {
    let normalizer = TextNormalizer::default();
    // The regex \b\d+\b doesn't match negative numbers (-5)
    // We'd need separate handling for negatives, so just test positive numbers
    let result = normalizer.normalize("Temperature is 5 degrees.").unwrap();
    assert!(result.contains("five degrees"));
}

#[test]
fn test_prosody_emphasis() {
    let config = ProsodyConfig {
        detect_emphasis: true,
        ..Default::default()
    };
    let prosody = Prosody::new(config);

    let result = prosody.predict("THIS IS IMPORTANT").unwrap();
    assert!(!result.is_empty());

    // Should detect emphasis from all caps
    let has_emphasis = result.iter().any(|a| {
        a.markers
            .iter()
            .any(|m| matches!(m, ProsodyMarker::Emphasis))
    });
    assert!(has_emphasis);
}

#[test]
fn test_prosody_no_emphasis_detection() {
    let config = ProsodyConfig {
        detect_emphasis: false,
        ..Default::default()
    };
    let prosody = Prosody::new(config);

    let result = prosody.predict("THIS IS TEXT").unwrap();

    // Should not detect emphasis when disabled
    let has_emphasis = result.iter().any(|a| {
        a.markers
            .iter()
            .any(|m| matches!(m, ProsodyMarker::Emphasis))
    });
    assert!(!has_emphasis);
}

#[test]
fn test_complete_workflow() {
    // Simulate a complete TTS workflow
    let pipeline = TextPipeline::default();

    // Input text with various features
    let input = "Dr. Johnson said: \"The meeting is at 3:30 PM. \
                 We need $500 for supplies. \
                 Will everyone attend? \
                 I hope so!\"";

    let result = pipeline.process_full(input).unwrap();

    // Verify all stages completed
    assert_eq!(result.original, input);
    assert!(!result.normalized.is_empty());
    assert!(!result.phonemes.is_empty());
    assert!(!result.prosody_annotations.is_empty());

    // Verify specific transformations
    assert!(result.normalized.contains("doctor"));

    // Verify phoneme output
    let phoneme_str = result.phoneme_string(true);
    assert!(!phoneme_str.is_empty());

    // Verify prosody markers
    let has_question = result
        .prosody_annotations
        .iter()
        .any(|a| a.markers.iter().any(|m| matches!(m, ProsodyMarker::Rising)));
    assert!(has_question);
}

#[test]
fn test_phoneme_string_output() {
    let pipeline = TextPipeline::default();
    let result = pipeline.process_full("hello world").unwrap();

    let with_stress = result.phoneme_string(true);
    let without_stress = result.phoneme_string(false);

    assert!(!with_stress.is_empty());
    assert!(!without_stress.is_empty());

    // With stress should potentially have digit markers
    assert!(with_stress.contains("HH") || with_stress.contains("W"));
    assert!(without_stress.contains("HH") || without_stress.contains("W"));
}

#[test]
fn test_duration_factors() {
    let prosody = Prosody::default();
    let result = prosody
        .predict("Short. This is a much longer sentence with many words.")
        .unwrap();

    assert!(!result.is_empty());
    // Duration factors should be set
    for annotation in &result {
        assert!(annotation.duration_factor > 0.0);
    }
}