scripts/test_generation.py

"""
Test generation script for ACE-Step Custom Edition
"""

import sys
from pathlib import Path
import logging

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent))

from src import ACEStepEngine, load_config, setup_logging

logger = setup_logging()


def test_basic_generation():
    """Test basic music generation."""
    logger.info("\n" + "="*60)
    logger.info("Testing Basic Music Generation")
    logger.info("="*60)
    
    try:
        # Load config
        config = load_config()
        
        # Initialize engine
        logger.info("Initializing ACE-Step engine...")
        engine = ACEStepEngine(config)
        
        # Test generation
        logger.info("\nGenerating test audio...")
        audio_path = engine.generate(
            prompt="A cheerful pop song with piano and drums",
            lyrics=None,
            duration=10,  # Short test
            temperature=0.7,
            seed=42
        )
        
        logger.info(f"\n✅ Generation successful!")
        logger.info(f"Output: {audio_path}")
        
        return True
        
    except Exception as e:
        logger.error(f"\n❌ Generation failed: {e}")
        import traceback
        traceback.print_exc()
        return False


def test_timeline_workflow():
    """Test timeline-based generation."""
    logger.info("\n" + "="*60)
    logger.info("Testing Timeline Workflow")
    logger.info("="*60)
    
    try:
        from src import TimelineManager
        
        config = load_config()
        timeline = TimelineManager(config)
        
        logger.info("Creating timeline...")
        timeline_id = timeline.create_timeline()
        logger.info(f"Timeline ID: {timeline_id}")
        
        logger.info("\n✅ Timeline system operational!")
        return True
        
    except Exception as e:
        logger.error(f"\n❌ Timeline test failed: {e}")
        return False


def test_audio_processing():
    """Test audio processing utilities."""
    logger.info("\n" + "="*60)
    logger.info("Testing Audio Processing")
    logger.info("="*60)
    
    try:
        from src import AudioProcessor
        import numpy as np
        
        config = load_config()
        processor = AudioProcessor(config)
        
        # Create test audio
        test_audio = np.random.randn(2, 44100) * 0.1  # 1 second stereo
        
        # Test normalization
        normalized = processor.normalize_audio(test_audio)
        logger.info("✓ Normalization OK")
        
        # Test crossfade
        crossfaded = processor.crossfade(test_audio, test_audio, fade_duration=0.5)
        logger.info("✓ Crossfade OK")
        
        logger.info("\n✅ Audio processing operational!")
        return True
        
    except Exception as e:
        logger.error(f"\n❌ Audio processing test failed: {e}")
        return False


def main():
    """Run all tests."""
    logger.info("""
    ╔══════════════════════════════════════════════════════════╗
    ║         ACE-Step Custom Edition - Test Suite            ║
    ╚══════════════════════════════════════════════════════════╝
    """)
    
    results = {
        "Audio Processing": test_audio_processing(),
        "Timeline Workflow": test_timeline_workflow(),
        # "Basic Generation": test_basic_generation(),  # Requires model download
    }
    
    # Summary
    logger.info("\n" + "="*60)
    logger.info("Test Summary")
    logger.info("="*60)
    
    all_passed = True
    for test_name, passed in results.items():
        status = "✅ PASS" if passed else "❌ FAIL"
        logger.info(f"{status} - {test_name}")
        if not passed:
            all_passed = False
    
    logger.info("="*60)
    
    if all_passed:
        logger.info("\n✅ All tests passed!")
        logger.info("\nNote: Full generation test requires model download.")
        logger.info("Run: python scripts/download_model.py")
    else:
        logger.error("\n❌ Some tests failed. Check logs above.")
        sys.exit(1)


if __name__ == "__main__":
    main()