Kokoro-FastAPI/Test Threads.py

#!/usr/bin/env python3
# Compatible with both Windows and Linux
"""
Kokoro TTS Race Condition Test

This script creates multiple concurrent requests to a Kokoro TTS service
to reproduce a race condition where audio outputs don't match the requested text.
Each thread generates a simple numbered sentence, which should make mismatches
easy to identify through listening.

To run:
python kokoro_race_condition_test.py --threads 8 --iterations 5 --url http://localhost:8880
"""

import argparse
import base64
import concurrent.futures
import json
import os
import requests
import time
import wave
import sys
from pathlib import Path


def setup_args():
    """Parse command line arguments"""
    parser = argparse.ArgumentParser(description="Test Kokoro TTS for race conditions")
    parser.add_argument("--url", default="http://localhost:8880",
                      help="Base URL of the Kokoro TTS service")
    parser.add_argument("--threads", type=int, default=8,
                      help="Number of concurrent threads to use")
    parser.add_argument("--iterations", type=int, default=5,
                      help="Number of iterations per thread")
    parser.add_argument("--voice", default="af_heart",
                      help="Voice to use for TTS")
    parser.add_argument("--output-dir", default="./tts_test_output",
                      help="Directory to save output files")
    parser.add_argument("--debug", action="store_true",
                      help="Enable debug logging")
    return parser.parse_args()


def generate_test_sentence(thread_id, iteration):
    """Generate a simple test sentence with numbers to make mismatches easily identifiable"""
    return f"This is test sentence number {thread_id}-{iteration}. " \
           f"If you hear this sentence, you should hear the numbers {thread_id}-{iteration}."


def log_message(message, debug=False, is_error=False):
    """Log messages with timestamps"""
    timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
    prefix = "[ERROR]" if is_error else "[INFO]"
    if is_error or debug:
        print(f"{prefix} {timestamp} - {message}")
    sys.stdout.flush()  # Ensure logs are visible in Docker output


def request_tts(url, test_id, text, voice, output_dir, debug=False):
    """Request TTS from the Kokoro API and save the WAV output"""
    start_time = time.time()
    output_file = os.path.join(output_dir, f"test_{test_id}.wav")
    text_file = os.path.join(output_dir, f"test_{test_id}.txt")
    
    # Log output paths for debugging
    log_message(f"Thread {test_id}: Text will be saved to: {text_file}", debug)
    log_message(f"Thread {test_id}: Audio will be saved to: {output_file}", debug)
    
    # Save the text for later comparison
    try:
        with open(text_file, "w") as f:
            f.write(text)
        log_message(f"Thread {test_id}: Successfully saved text file", debug)
    except Exception as e:
        log_message(f"Thread {test_id}: Error saving text file: {str(e)}", debug, is_error=True)
    
    # Make the TTS request
    try:
        log_message(f"Thread {test_id}: Requesting TTS for: '{text}'", debug)
        
        response = requests.post(
            f"{url}/v1/audio/speech",
            json={
                "model": "kokoro",
                "input": text,
                "voice": voice,
                "response_format": "wav"
            },
            headers={"Accept": "audio/wav"},
            timeout=60  # Increase timeout to 60 seconds
        )
        
        log_message(f"Thread {test_id}: Response status code: {response.status_code}", debug)
        log_message(f"Thread {test_id}: Response content type: {response.headers.get('Content-Type', 'None')}", debug)
        log_message(f"Thread {test_id}: Response content length: {len(response.content)} bytes", debug)
        
        if response.status_code != 200:
            log_message(f"Thread {test_id}: API error: {response.status_code} - {response.text}", debug, is_error=True)
            return False
        
        # Check if we got valid audio data
        if len(response.content) < 100:  # Sanity check - WAV files should be larger than this
            log_message(f"Thread {test_id}: Received suspiciously small audio data: {len(response.content)} bytes", debug, is_error=True)
            log_message(f"Thread {test_id}: Content (base64): {base64.b64encode(response.content).decode('utf-8')}", debug, is_error=True)
            return False
        
        # Save the audio output with explicit error handling
        try:
            with open(output_file, "wb") as f:
                bytes_written = f.write(response.content)
                log_message(f"Thread {test_id}: Wrote {bytes_written} bytes to {output_file}", debug)
            
            # Verify the WAV file exists and has content
            if os.path.exists(output_file):
                file_size = os.path.getsize(output_file)
                log_message(f"Thread {test_id}: Verified file exists with size: {file_size} bytes", debug)
                
                # Validate WAV file by reading its headers
                try:
                    with wave.open(output_file, 'rb') as wav_file:
                        channels = wav_file.getnchannels()
                        sample_width = wav_file.getsampwidth()
                        framerate = wav_file.getframerate()
                        frames = wav_file.getnframes()
                        log_message(f"Thread {test_id}: Valid WAV file - channels: {channels}, "
                                   f"sample width: {sample_width}, framerate: {framerate}, frames: {frames}", debug)
                except Exception as wav_error:
                    log_message(f"Thread {test_id}: Invalid WAV file: {str(wav_error)}", debug, is_error=True)
            else:
                log_message(f"Thread {test_id}: File was not created: {output_file}", debug, is_error=True)
        except Exception as save_error:
            log_message(f"Thread {test_id}: Error saving audio file: {str(save_error)}", debug, is_error=True)
            return False
        
        end_time = time.time()
        log_message(f"Thread {test_id}: Saved output to {output_file} (time: {end_time - start_time:.2f}s)", debug)
        return True
    
    except requests.exceptions.Timeout:
        log_message(f"Thread {test_id}: Request timed out", debug, is_error=True)
        return False
    except Exception as e:
        log_message(f"Thread {test_id}: Exception: {str(e)}", debug, is_error=True)
        return False


def worker_task(thread_id, args):
    """Worker task for each thread"""
    for i in range(args.iterations):
        iteration = i + 1
        test_id = f"{thread_id:02d}_{iteration:02d}"
        text = generate_test_sentence(thread_id, iteration)
        success = request_tts(args.url, test_id, text, args.voice, args.output_dir, args.debug)
        
        if not success:
            log_message(f"Thread {thread_id}: Iteration {iteration} failed", args.debug, is_error=True)
        
        # Small delay between iterations to avoid overwhelming the API
        time.sleep(0.1)


def run_test(args):
    """Run the test with the specified parameters"""
    # Ensure output directory exists and check permissions
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Test write access to the output directory
    test_file = os.path.join(args.output_dir, "write_test.txt")
    try:
        with open(test_file, "w") as f:
            f.write("Testing write access\n")
        os.remove(test_file)
        log_message(f"Successfully verified write access to output directory: {args.output_dir}")
    except Exception as e:
        log_message(f"Warning: Cannot write to output directory {args.output_dir}: {str(e)}", is_error=True)
        log_message(f"Current directory: {os.getcwd()}", is_error=True)
        log_message(f"Directory contents: {os.listdir('.')}", is_error=True)
    
    # Test connection to Kokoro TTS service
    try:
        response = requests.get(f"{args.url}/health", timeout=5)
        if response.status_code == 200:
            log_message(f"Successfully connected to Kokoro TTS service at {args.url}")
        else:
            log_message(f"Warning: Kokoro TTS service health check returned status {response.status_code}", is_error=True)
    except Exception as e:
        log_message(f"Warning: Cannot connect to Kokoro TTS service at {args.url}: {str(e)}", is_error=True)
    
    # Record start time
    start_time = time.time()
    log_message(f"Starting test with {args.threads} threads, {args.iterations} iterations per thread")
    
    # Create and start worker threads
    with concurrent.futures.ThreadPoolExecutor(max_workers=args.threads) as executor:
        futures = []
        for thread_id in range(1, args.threads + 1):
            futures.append(executor.submit(worker_task, thread_id, args))
        
        # Wait for all tasks to complete
        for future in concurrent.futures.as_completed(futures):
            try:
                future.result()
            except Exception as e:
                log_message(f"Thread execution failed: {str(e)}", args.debug, is_error=True)
    
    # Record end time and print summary
    end_time = time.time()
    total_time = end_time - start_time
    total_requests = args.threads * args.iterations
    log_message(f"Test completed in {total_time:.2f} seconds")
    log_message(f"Total requests: {total_requests}")
    log_message(f"Average time per request: {total_time / total_requests:.2f} seconds")
    log_message(f"Requests per second: {total_requests / total_time:.2f}")
    log_message(f"Output files saved to: {os.path.abspath(args.output_dir)}")
    log_message("To verify, listen to the audio files and check if they match the text files")
    log_message("If you hear audio describing a different test number than the filename, you've found a race condition")


def analyze_audio_files(output_dir):
    """Provide summary of the generated audio files"""
    # Look for both WAV and TXT files
    wav_files = list(Path(output_dir).glob("*.wav"))
    txt_files = list(Path(output_dir).glob("*.txt"))
    
    log_message(f"Found {len(wav_files)} WAV files and {len(txt_files)} TXT files")
    
    if len(wav_files) == 0:
        log_message("No WAV files found! This indicates the TTS service requests may be failing.", is_error=True)
        log_message("Check the connection to the TTS service and the response status codes above.", is_error=True)
    
    file_stats = []
    for wav_path in wav_files:
        try:
            with wave.open(str(wav_path), 'rb') as wav_file:
                frames = wav_file.getnframes()
                rate = wav_file.getframerate()
                duration = frames / rate
                
                # Get corresponding text
                text_path = wav_path.with_suffix('.txt')
                if text_path.exists():
                    with open(text_path, 'r') as text_file:
                        text = text_file.read().strip()
                else:
                    text = "N/A"
                
                file_stats.append({
                    'filename': wav_path.name,
                    'duration': duration,
                    'text': text
                })
        except Exception as e:
            log_message(f"Error analyzing {wav_path}: {str(e)}", False, is_error=True)
    
    # Print summary table
    if file_stats:
        log_message("\nAudio File Summary:")
        log_message(f"{'Filename':<20}{'Duration':<12}{'Text':<60}")
        log_message("-" * 92)
        for stat in file_stats:
            log_message(f"{stat['filename']:<20}{stat['duration']:<12.2f}{stat['text'][:57]+'...' if len(stat['text']) > 60 else stat['text']:<60}")
    
    # List missing WAV files where text files exist
    missing_wavs = set(p.stem for p in txt_files) - set(p.stem for p in wav_files)
    if missing_wavs:
        log_message(f"\nFound {len(missing_wavs)} text files without corresponding WAV files:", is_error=True)
        for stem in sorted(list(missing_wavs))[:10]:  # Limit to 10 for readability
            log_message(f"  - {stem}.txt (no WAV file)", is_error=True)
        if len(missing_wavs) > 10:
            log_message(f"  ... and {len(missing_wavs) - 10} more", is_error=True)


if __name__ == "__main__":
    args = setup_args()
    run_test(args)
    analyze_audio_files(args.output_dir)
    
    log_message("\nNext Steps:")
    log_message("1. Listen to the generated audio files")
    log_message("2. Verify if each audio correctly says its ID number")
    log_message("3. Check for any mismatches between the audio content and the text files")
    log_message("4. If mismatches are found, you've successfully reproduced the race condition")