audiobook/backend/test_segment_based_generation.py

"""
Test for segment-based audiobook generation (Task 19).

Tests that the TTS endpoint can accept edited script segments
and generate audio from them, while maintaining backward compatibility
with text-based generation.
"""
import pytest
from fastapi.testclient import TestClient
from main import app

client = TestClient(app)


def test_generate_audiobook_with_segments():
    """Test that audiobook can be generated from edited script segments."""
    # Prepare request with segments
    request_data = {
        "segments": [
            {"id": "seg_0", "speaker": "Alice", "text": "Hello, how are you?"},
            {"id": "seg_1", "speaker": "Bob", "text": "I'm doing great, thanks!"},
            {"id": "seg_2", "speaker": "Narrator", "text": "They continued their conversation."}
        ],
        "lang": "a",
        "roles": [
            {"id": "Alice", "voice": "af_bella"},
            {"id": "Bob", "voice": "am_adam"},
            {"id": "Narrator", "voice": "af_sarah"}
        ],
        "speed": 1.0
    }
    
    response = client.post("/api/tts", json=request_data)
    
    # Should succeed
    assert response.status_code == 200
    assert response.headers["content-type"] == "audio/wav"
    assert "audiobook.wav" in response.headers.get("content-disposition", "")
    assert len(response.content) > 0


def test_generate_audiobook_with_text_backward_compatibility():
    """Test that audiobook generation still works with text (backward compatibility)."""
    # Prepare request with text (old format)
    request_data = {
        "text": "Alice: Hello there!\nBob: Hi Alice!",
        "lang": "a",
        "roles": [
            {"id": "Alice", "voice": "af_bella"},
            {"id": "Bob", "voice": "am_adam"}
        ],
        "speed": 1.0
    }
    
    response = client.post("/api/tts", json=request_data)
    
    # Should succeed
    assert response.status_code == 200
    assert response.headers["content-type"] == "audio/wav"
    assert len(response.content) > 0


def test_generate_audiobook_segments_take_precedence():
    """Test that segments are used when both segments and text are provided."""
    # Prepare request with both segments and text
    request_data = {
        "segments": [
            {"id": "seg_0", "speaker": "Alice", "text": "From segments"}
        ],
        "text": "Alice: From text",
        "lang": "a",
        "roles": [
            {"id": "Alice", "voice": "af_bella"}
        ],
        "speed": 1.0
    }
    
    response = client.post("/api/tts", json=request_data)
    
    # Should succeed and use segments
    assert response.status_code == 200
    assert response.headers["content-type"] == "audio/wav"


def test_generate_audiobook_missing_both_segments_and_text():
    """Test that request fails when neither segments nor text is provided."""
    request_data = {
        "lang": "a",
        "roles": [
            {"id": "Alice", "voice": "af_bella"}
        ],
        "speed": 1.0
    }
    
    response = client.post("/api/tts", json=request_data)
    
    # Should fail with validation error
    assert response.status_code == 400
    assert "ValidationError" in response.json()["error"]


def test_generate_audiobook_empty_segments():
    """Test that request fails when segments array is empty."""
    request_data = {
        "segments": [],
        "lang": "a",
        "roles": [
            {"id": "Alice", "voice": "af_bella"}
        ],
        "speed": 1.0
    }
    
    response = client.post("/api/tts", json=request_data)
    
    # Should fail - empty segments should be treated as not provided
    assert response.status_code == 400


def test_edited_script_preserves_speaker_assignments():
    """Test that edited speaker assignments in segments are preserved."""
    # Segments where user has changed speaker assignments
    request_data = {
        "segments": [
            {"id": "seg_0", "speaker": "Bob", "text": "This was originally Alice's line"},
            {"id": "seg_1", "speaker": "Alice", "text": "This was originally Bob's line"}
        ],
        "lang": "a",
        "roles": [
            {"id": "Alice", "voice": "af_bella"},
            {"id": "Bob", "voice": "am_adam"}
        ],
        "speed": 1.0
    }
    
    response = client.post("/api/tts", json=request_data)
    
    # Should succeed with edited assignments
    assert response.status_code == 200
    assert response.headers["content-type"] == "audio/wav"
    assert len(response.content) > 0


def test_segment_with_unmapped_speaker_uses_default():
    """Test that segments with unmapped speakers fall back to default voice."""
    request_data = {
        "segments": [
            {"id": "seg_0", "speaker": "UnknownSpeaker", "text": "Hello from unknown"}
        ],
        "lang": "a",
        "roles": [
            {"id": "Alice", "voice": "af_bella"}
        ],
        "speed": 1.0
    }
    
    response = client.post("/api/tts", json=request_data)
    
    # Should succeed using default voice
    assert response.status_code == 200
    assert response.headers["content-type"] == "audio/wav"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])