#!/usr/bin/env python3 """ Unit tests for vocal captions and pitch range integration using unittest framework. Tests compatibility with: 1. Hook mode + Add vocal 2. Mumble mode + Add vocal 3. Vox block + Add Vocal 4. Regular mode with vocal captions """ import unittest import sys import random from pathlib import Path # Add sunoGPT to path sys.path.insert(0, str(Path(__file__).parent.parent)) from data_utils import extract_vocal_captions from text_utils import build_text, get_control_tags class TestVocalCaptionPitchRange(unittest.TestCase): """Unit tests for vocal captions and pitch range integration.""" def setUp(self): """Set up test fixtures.""" # Set random seed for reproducible tests random.seed(42) self.test_metadata = { "id": "test_vocal_integration", "tags": ["pop", "upbeat", "modern"], "text": "This is a test song with vocal content", "stems_captions": { "Vocals": [ { "prompt_type": "voice_description_keywords", "caption": "Male, adult, smooth, clear, resonant, soulful, expressive, warm vocal", "error": None, } ], "Backing_Vocals": [ { "prompt_type": "voice_description_keywords", "caption": "Female, harmony, soft, melodic voice", "error": None, } ], "Bass": [ { "prompt_type": "musical_role", "caption": "Deep, rhythmic, foundational bass", "error": None, } ], }, "vocal_pitch_range": {"IQR_1.5x": {"min_hz": 85.2, "max_hz": 350.8}}, } def test_extract_vocal_captions(self): """Test the extract_vocal_captions function.""" vocal_tags = extract_vocal_captions(self.test_metadata) self.assertIsNotNone(vocal_tags, "Vocal captions should be extracted") self.assertIsInstance(vocal_tags, list, "Vocal tags should be a list") self.assertGreater(len(vocal_tags), 0, "Should extract vocal keywords") # Check that only Vocals stem keywords are extracted (not Backing_Vocals due to current logic) expected_keywords = [ "Male", "adult", "smooth", "clear", "resonant", "soulful", "expressive", "warm vocal", ] self.assertEqual( vocal_tags, expected_keywords, f"Expected {expected_keywords} but got {vocal_tags}" ) def test_extract_vocal_captions_no_data(self): """Test extract_vocal_captions with no stems_captions data.""" empty_metadata = {"id": "test", "tags": ["pop"]} vocal_tags = extract_vocal_captions(empty_metadata) self.assertIsNone(vocal_tags, "Should return None when no stems_captions") def test_extract_vocal_captions_no_vocal_stems(self): """Test extract_vocal_captions with no vocal stems.""" metadata_no_vocals = { "stems_captions": { "Bass": [{"prompt_type": "voice_description_keywords", "caption": "Deep bass"}], "Drums": [{"prompt_type": "voice_description_keywords", "caption": "Heavy drums"}], } } vocal_tags = extract_vocal_captions(metadata_no_vocals) self.assertIsNone(vocal_tags, "Should return None when no vocal stems") def test_build_text_vocal_integration_inference(self): """Test build_text with vocal tags in inference mode.""" vocal_tags = extract_vocal_captions(self.test_metadata) result_text = build_text( tags=["pop", "upbeat", "modern"], text="Test lyrics for vocal integration", sample_duration_s=30.0, sample_duration_toks=750, vocal_tags=vocal_tags, vocal_pitch_hz_min=85.2, vocal_pitch_hz_max=350.8, inference=True, ) self.assertIsInstance(result_text, str, "Should return string") self.assertIn("vocal", result_text.lower(), "Should contain vocal-related tags") self.assertIn("vocal_pitch_hz", result_text, "Should contain pitch range tags") # Check that all expected vocal tags are present in the result expected_vocal_tags = [ "Male", "adult", "smooth", "clear", "resonant", "soulful", "expressive", "warm vocal", ] result_lower = result_text.lower() vocal_tags_found = [tag for tag in expected_vocal_tags if tag.lower() in result_lower] self.assertTrue( all(tag.lower() in result_lower for tag in expected_vocal_tags), f"Should contain all expected vocal tags, found: {vocal_tags_found}, missing: {set(expected_vocal_tags) - set(vocal_tags_found)}", ) def test_build_text_vocal_integration_training(self): """Test build_text with vocal tags in training mode.""" vocal_tags = extract_vocal_captions(self.test_metadata) result_text = build_text( tags=["pop", "upbeat", "modern"], text="Test lyrics for vocal integration", sample_duration_s=30.0, sample_duration_toks=750, vocal_tags=vocal_tags, vocal_pitch_hz_min=85.2, vocal_pitch_hz_max=350.8, inference=True, # Use inference mode to ensure control tags are included ) self.assertIsInstance(result_text, str, "Should return string") self.assertIn("vocal", result_text.lower(), "Should contain vocal-related tags") self.assertIn("vocal_pitch_hz", result_text, "Should contain pitch range tags") def test_build_text_no_vocal_tags(self): """Test build_text with no vocal tags.""" result_text = build_text( tags=["pop", "upbeat"], text="Test lyrics", sample_duration_s=25.0, sample_duration_toks=625, vocal_tags=None, inference=True, # Use inference to ensure tags are included ) self.assertIsInstance(result_text, str, "Should return string") # Should still work without vocal tags self.assertIn("pop", result_text.lower(), "Should contain regular tags") def test_mode_compatibility_hook(self): """Test compatibility with hook mode.""" vocal_tags = extract_vocal_captions(self.test_metadata) result_text = build_text( tags=["test", "compatibility"], text="Test mode compatibility", sample_duration_s=25.0, sample_duration_toks=625, vocal_tags=vocal_tags, vocal_pitch_hz_min=85.2, vocal_pitch_hz_max=350.8, hook_only=True, inference=False, ) self.assertIsInstance(result_text, str, "Hook mode should work with vocal integration") def test_mode_compatibility_chorus(self): """Test compatibility with chorus mode.""" vocal_tags = extract_vocal_captions(self.test_metadata) result_text = build_text( tags=["test", "compatibility"], text="Test mode compatibility", sample_duration_s=25.0, sample_duration_toks=625, vocal_tags=vocal_tags, vocal_pitch_hz_min=85.2, vocal_pitch_hz_max=350.8, start_from_chorus=True, inference=False, ) self.assertIsInstance(result_text, str, "Chorus mode should work with vocal integration") def test_mode_compatibility_regular(self): """Test compatibility with regular mode.""" vocal_tags = extract_vocal_captions(self.test_metadata) result_text = build_text( tags=["test", "compatibility"], text="Test mode compatibility", sample_duration_s=25.0, sample_duration_toks=625, vocal_tags=vocal_tags, vocal_pitch_hz_min=85.2, vocal_pitch_hz_max=350.8, hook_only=False, start_from_chorus=False, inference=False, ) self.assertIsInstance(result_text, str, "Regular mode should work with vocal integration") def test_get_control_tags_with_pitch_range(self): """Test get_control_tags with vocal pitch range.""" control_tags = get_control_tags( sample_duration_s=30.0, sample_duration_toks=750, vocal_pitch_hz_min=85.2, vocal_pitch_hz_max=350.8, do_augment=False, # Disable augmentation to ensure pitch tags are included ) self.assertIsNotNone(control_tags, "Should return control tags") self.assertIn("vocal_pitch_hz_min:85", control_tags, "Should include min pitch") self.assertIn("vocal_pitch_hz_max:351", control_tags, "Should include max pitch") def test_get_control_tags_without_pitch_range(self): """Test get_control_tags without vocal pitch range.""" control_tags = get_control_tags( sample_duration_s=30.0, sample_duration_toks=750, vocal_pitch_hz_min=None, vocal_pitch_hz_max=None, do_augment=False, # Disable augmentation to ensure we get control tags ) self.assertIsNotNone(control_tags, "Should return control tags") self.assertNotIn("vocal_pitch_hz", control_tags, "Should not include pitch tags") class TestVocalCaptionPitchRangeEdgeCases(unittest.TestCase): """Test edge cases for vocal captions and pitch range integration.""" def test_extract_vocal_captions_empty_caption(self): """Test extraction with empty caption text.""" metadata = { "stems_captions": { "Vocals": [{"prompt_type": "voice_description_keywords", "caption": "", "error": None}] } } vocal_tags = extract_vocal_captions(metadata) self.assertIsNone(vocal_tags, "Should return None for empty captions") def test_extract_vocal_captions_wrong_prompt_type(self): """Test extraction with wrong prompt type.""" metadata = { "stems_captions": { "Vocals": [{"prompt_type": "musical_role", "caption": "Lead vocals", "error": None}] } } vocal_tags = extract_vocal_captions(metadata) self.assertIsNone(vocal_tags, "Should return None for wrong prompt type") def test_build_text_empty_vocal_tags(self): """Test build_text with empty vocal tags list.""" result_text = build_text( tags=["pop"], text="Test", sample_duration_s=20.0, sample_duration_toks=500, vocal_tags=[], inference=False, ) self.assertIsInstance(result_text, str, "Should handle empty vocal tags list") class TestVocalCaptionProcessingPaths(unittest.TestCase): """Test the distribution and behavior of vocal caption processing paths.""" def setUp(self): """Set up test fixtures with random seed for reproducible path testing.""" # Set random seed for reproducible tests random.seed(123) def test_vocal_paths_distribution(self): """Test that both vocal tag processing paths work multiple times.""" vocal_tags = ["Male", "smooth", "clear", "resonant"] regular_tags = ["pop", "upbeat"] path_1_count = 0 # Merged tags path_2_count = 0 # Separate [vocal:...] # Run multiple times to test both paths for _ in range(20): # Increased iterations for better statistics result = build_text( tags=regular_tags, text="Test lyrics", sample_duration_s=30.0, sample_duration_toks=750, vocal_tags=vocal_tags, inference=False, ) has_vocal_element = "[vocal:" in result has_merged_tags = not has_vocal_element and any( word in result.lower() for word in ["male", "smooth", "clear", "resonant"] ) if has_vocal_element: path_2_count += 1 elif has_merged_tags: path_1_count += 1 # Verify both paths are working (with some tolerance for randomness) self.assertGreater(path_1_count, 0, "Should have some merged tags processing") self.assertGreater(path_2_count, 0, "Should have some separate [vocal:] processing") # At least one of each path should occur in 20 runs (very high probability) total_vocal_processing = path_1_count + path_2_count self.assertGreater( total_vocal_processing, 10, f"Expected most runs to process vocal tags, got {total_vocal_processing}/20", ) def test_path_1_merged_tags(self): """Test Path 1 - merged tags with vocal suffixes.""" vocal_tags = ["Male", "smooth"] regular_tags = ["pop", "upbeat"] # Run until we get a Path 1 result (merged tags) for _ in range(50): # Max attempts to avoid infinite loop result = build_text( tags=regular_tags, text="Test lyrics", sample_duration_s=30.0, sample_duration_toks=750, vocal_tags=vocal_tags, inference=False, ) has_vocal_element = "[vocal:" in result has_merged_tags = not has_vocal_element and any( word in result.lower() for word in ["male", "smooth"] ) if has_merged_tags: # Verify it's a merged result (no separate [vocal:] section) self.assertNotIn("[vocal:", result, "Path 1 should not have separate vocal section") # Verify it contains both regular and vocal tags result_lower = result.lower() self.assertTrue( any(tag.lower() in result_lower for tag in regular_tags), "Should contain regular tags", ) self.assertTrue( any(tag.lower() in result_lower for tag in vocal_tags), "Should contain vocal tags" ) return # Success - found Path 1 result self.fail("Could not generate Path 1 (merged) result in 50 attempts") def test_path_2_separate_vocal_element(self): """Test Path 2 - separate [vocal:...] element.""" vocal_tags = ["Male", "smooth"] regular_tags = ["pop", "upbeat"] # Run until we get a Path 2 result (separate vocal element) for _ in range(50): # Max attempts to avoid infinite loop result = build_text( tags=regular_tags, text="Test lyrics", sample_duration_s=30.0, sample_duration_toks=750, vocal_tags=vocal_tags, inference=False, ) if "[vocal:" in result: # Verify it has separate vocal section self.assertIn("[vocal:", result, "Path 2 should have separate vocal section") return # Success - found Path 2 result self.fail("Could not generate Path 2 (separate vocal) result in 50 attempts") if __name__ == "__main__": unittest.main()