"""Comprehensive tests for seed_utils module.""" import pytest import numpy as np from unittest.mock import patch from suno_recs.worker.retrievals.seed_utils import ( weighted_sample_without_replacement, get_video_knn_seeds, get_audio_knn_seeds, ) from suno_recs.worker.constants import RECS_CONFIG class TestWeightedSampleWithoutReplacement: """Test suite for weighted_sample_without_replacement function.""" def test_basic_sampling(self): """Test basic weighted sampling functionality.""" categories = [ {"ids": ["a", "b", "c"], "weight": 2.0, "name": "cat1"}, {"ids": ["d", "e", "f"], "weight": 1.0, "name": "cat2"}, ] result = weighted_sample_without_replacement(categories, n=4, seed=42) # Should return 4 items assert len(result) == 4 # All items should be unique assert len(set(result)) == 4 # All items should be from the original sets assert all(item in ["a", "b", "c", "d", "e", "f"] for item in result) def test_deterministic_with_seed(self): """Test that same seed produces same results.""" categories = [ {"ids": list("abcdef"), "weight": 1.0}, {"ids": list("ghijkl"), "weight": 1.0}, ] result1 = weighted_sample_without_replacement(categories, n=6, seed=123) result2 = weighted_sample_without_replacement(categories, n=6, seed=123) assert result1 == result2 def test_different_seeds_different_results(self): """Test that different seeds produce different results.""" categories = [ {"ids": list("abcdef"), "weight": 1.0}, {"ids": list("ghijkl"), "weight": 1.0}, ] result1 = weighted_sample_without_replacement(categories, n=6, seed=123) result2 = weighted_sample_without_replacement(categories, n=6, seed=456) # Very unlikely to be the same with different seeds assert result1 != result2 def test_weight_bias(self): """Test that higher weights lead to more selections.""" # Category 1 has 10x weight of category 2 categories = [ {"ids": ["a"] * 100, "weight": 10.0, "name": "high_weight"}, {"ids": ["b"] * 100, "weight": 1.0, "name": "low_weight"}, ] result = weighted_sample_without_replacement(categories, n=100, seed=42) a_count = result.count("a") b_count = result.count("b") # With 10:1 weight ratio, we expect roughly 10x more 'a' than 'b' # Allow some variance due to randomness assert a_count > b_count * 5 # At least 5x more assert a_count + b_count == 100 def test_exhausted_category(self): """Test behavior when a category runs out of items.""" categories = [ {"ids": ["a", "b"], "weight": 10.0}, # Only 2 items but high weight {"ids": ["c", "d", "e", "f", "g"], "weight": 1.0}, ] result = weighted_sample_without_replacement(categories, n=7, seed=42) # Should get all 7 items even though first category only has 2 assert len(result) == 7 assert set(result) == {"a", "b", "c", "d", "e", "f", "g"} def test_request_more_than_available(self): """Test requesting more items than available.""" categories = [ {"ids": ["a", "b"], "weight": 1.0}, {"ids": ["c", "d"], "weight": 1.0}, ] result = weighted_sample_without_replacement(categories, n=10, seed=42) # Should only get 4 items (all available) assert len(result) == 4 assert set(result) == {"a", "b", "c", "d"} def test_empty_categories(self): """Test with empty categories.""" categories = [ {"ids": [], "weight": 1.0}, {"ids": ["a", "b"], "weight": 1.0}, ] result = weighted_sample_without_replacement(categories, n=3, seed=42) # Should only sample from non-empty category assert len(result) == 2 assert set(result) == {"a", "b"} def test_zero_weight(self): """Test categories with zero weight are ignored.""" categories = [ {"ids": ["a", "b"], "weight": 0.0}, {"ids": ["c", "d"], "weight": 1.0}, ] result = weighted_sample_without_replacement(categories, n=3, seed=42) # Should only sample from positive weight category assert len(result) == 2 assert set(result) == {"c", "d"} def test_all_empty_or_zero_weight(self): """Test when all categories are empty or have zero weight.""" categories = [ {"ids": [], "weight": 1.0}, {"ids": ["a"], "weight": 0.0}, ] result = weighted_sample_without_replacement(categories, n=5, seed=42) assert result == [] class TestGetVideoKnnSeeds: """Test suite for get_video_knn_seeds function.""" def test_basic_video_seeds(self): """Test basic video seed selection.""" liked_hooks = ["h1", "h2", "h3", "h4", "h5"] positive_hooks = ["p1", "p2", "p3"] seeds, breakdown = get_video_knn_seeds(liked_hooks, positive_hooks) # Check we don't exceed max assert len(seeds) <= RECS_CONFIG.get("max_video_knn_hook_seeds", 30) # Check breakdown assert breakdown["total"] == len(seeds) assert breakdown["liked_hooks"] + breakdown["positive_interaction"] == breakdown["total"] # All seeds should be from input lists all_hooks = set(liked_hooks + positive_hooks) assert all(s in all_hooks for s in seeds) def test_video_seeds_weight_ratio(self): """Test that video seeds respect weight ratios.""" # Default is 2:1 ratio (liked:positive) liked_hooks = [f"l{i}" for i in range(100)] positive_hooks = [f"p{i}" for i in range(100)] # Run multiple times and check average ratio liked_counts = [] positive_counts = [] for seed in range(10): seeds, breakdown = get_video_knn_seeds(liked_hooks, positive_hooks) liked_counts.append(breakdown["liked_hooks"]) positive_counts.append(breakdown["positive_interaction"]) avg_liked = np.mean(liked_counts) avg_positive = np.mean(positive_counts) # Should be roughly 2:1 ratio ratio = avg_liked / avg_positive assert 1.5 < ratio < 2.5 # Allow some variance def test_video_seeds_empty_inputs(self): """Test video seeds with empty inputs.""" # Empty liked hooks seeds1, breakdown1 = get_video_knn_seeds([], ["p1", "p2"]) assert len(seeds1) == 2 assert breakdown1["liked_hooks"] == 0 assert breakdown1["positive_interaction"] == 2 # Empty positive hooks seeds2, breakdown2 = get_video_knn_seeds(["l1", "l2"], []) assert len(seeds2) == 2 assert breakdown2["liked_hooks"] == 2 assert breakdown2["positive_interaction"] == 0 # Both empty seeds3, breakdown3 = get_video_knn_seeds([], []) assert len(seeds3) == 0 assert breakdown3["total"] == 0 def test_video_seeds_max_limit(self): """Test video seeds respect max limit.""" liked_hooks = [f"l{i}" for i in range(100)] positive_hooks = [f"p{i}" for i in range(100)] seeds, _ = get_video_knn_seeds(liked_hooks, positive_hooks) max_seeds = RECS_CONFIG.get("max_video_knn_hook_seeds", 30) assert len(seeds) <= max_seeds def test_video_seeds_with_custom_config(self): """Test video seeds with custom boost configuration.""" with patch.dict(RECS_CONFIG, { "video_seed_boosts": {"liked_hooks": 1, "positive_hooks": 3}, "max_video_knn_hook_seeds": 10 }): liked_hooks = [f"l{i}" for i in range(20)] positive_hooks = [f"p{i}" for i in range(20)] seeds, breakdown = get_video_knn_seeds(liked_hooks, positive_hooks) assert len(seeds) <= 10 # With 1:3 ratio, should have more positive than liked assert breakdown["positive_interaction"] > breakdown["liked_hooks"] class TestGetAudioKnnSeeds: """Test suite for get_audio_knn_seeds function.""" def test_basic_audio_seeds(self): """Test basic audio seed selection.""" liked_hooks = ["lh1", "lh2", "lh3"] positive_hooks = ["ph1", "ph2"] liked_clips = ["lc1", "lc2", "lc3", "lc4"] listening_clips = ["ls1", "ls2"] hook_like, hook_pos, clips, details = get_audio_knn_seeds( liked_hooks, positive_hooks, liked_clips, listening_clips ) # Check total doesn't exceed max total = len(hook_like) + len(hook_pos) + len(clips) assert total <= RECS_CONFIG.get("max_audio_knn_seeds", 40) assert details["total"] == total # Check seeds are properly categorized assert all(s in liked_hooks for s in hook_like) assert all(s in positive_hooks for s in hook_pos) assert all(s in (liked_clips + listening_clips) for s in clips) # Check details structure assert "available" in details assert "boosts" in details assert "selected_counts" in details def test_audio_seeds_weight_distribution(self): """Test audio seeds respect weight distribution.""" # Default weights: liked_hooks(3), positive_hooks(1), liked_clips(2), listening_clips(1) liked_hooks = [f"lh{i}" for i in range(50)] positive_hooks = [f"ph{i}" for i in range(50)] liked_clips = [f"lc{i}" for i in range(50)] listening_clips = [f"ls{i}" for i in range(50)] # Run multiple times to check distribution counts = {"liked_hooks": 0, "positive_hooks": 0, "liked_clips": 0, "listening_clips": 0} for seed in range(10): _, _, _, details = get_audio_knn_seeds( liked_hooks, positive_hooks, liked_clips, listening_clips ) for key in counts: counts[key] += details["selected_counts"][key] # Check ratios roughly match weights (3:1:2:1) # Normalize by the smallest count min_count = min(counts.values()) if min_count > 0: ratios = {k: v / min_count for k, v in counts.items()} # Allow some variance due to randomness assert 1.5 < ratios["liked_hooks"] < 4.5 # Should be ~3 assert 0.3 < ratios["positive_hooks"] < 2.0 # Should be ~1 assert 1.0 < ratios["liked_clips"] < 3.0 # Should be ~2 assert 0.3 < ratios["listening_clips"] < 2.0 # Should be ~1 def test_audio_seeds_empty_categories(self): """Test audio seeds with empty categories.""" # Only liked hooks hook_like, hook_pos, clips, details = get_audio_knn_seeds( ["lh1", "lh2"], [], [], [] ) assert len(hook_like) == 2 assert len(hook_pos) == 0 assert len(clips) == 0 # Only clips hook_like, hook_pos, clips, details = get_audio_knn_seeds( [], [], ["lc1", "lc2"], ["ls1"] ) assert len(hook_like) == 0 assert len(hook_pos) == 0 assert len(clips) == 3 # All empty hook_like, hook_pos, clips, details = get_audio_knn_seeds([], [], [], []) assert details["total"] == 0 def test_audio_seeds_max_limit(self): """Test audio seeds respect max limit.""" # Create many seeds liked_hooks = [f"lh{i}" for i in range(100)] positive_hooks = [f"ph{i}" for i in range(100)] liked_clips = [f"lc{i}" for i in range(100)] listening_clips = [f"ls{i}" for i in range(100)] hook_like, hook_pos, clips, details = get_audio_knn_seeds( liked_hooks, positive_hooks, liked_clips, listening_clips ) total = len(hook_like) + len(hook_pos) + len(clips) max_seeds = RECS_CONFIG.get("max_audio_knn_seeds", 40) assert total <= max_seeds assert details["total"] == total def test_audio_seeds_available_counts(self): """Test that available counts are reported correctly.""" liked_hooks = ["lh1", "lh2", "lh3"] positive_hooks = ["ph1"] liked_clips = ["lc1", "lc2"] listening_clips = ["ls1", "ls2", "ls3", "ls4"] _, _, _, details = get_audio_knn_seeds( liked_hooks, positive_hooks, liked_clips, listening_clips ) assert details["available"]["liked_hooks"] == 3 assert details["available"]["positive_hooks"] == 1 assert details["available"]["liked_clips"] == 2 assert details["available"]["listening_clips"] == 4 def test_audio_seeds_overlapping_inputs(self): """Test audio seeds behavior with overlapping inputs.""" # Create overlapping inputs liked_hooks = ["h1", "h2", "h3"] positive_hooks = ["h2", "h3", "h4"] # h2 and h3 overlap with liked hook_like, hook_pos, clips, details = get_audio_knn_seeds( liked_hooks, positive_hooks, [], [] ) # When IDs overlap between categories, they can be selected multiple times # from different categories during weighted sampling, and then appear as # duplicates when re-categorized based on set membership. # This is the expected behavior of the current implementation. # All selected hooks should be from the input sets all_hooks = set(liked_hooks + positive_hooks) assert all(h in all_hooks for h in hook_like) assert all(h in all_hooks for h in hook_pos) # The total in details represents the number of items sampled by # weighted_sample_without_replacement, not the sum of categorized lists # (which can have duplicates when IDs overlap between categories) assert details["total"] <= len(hook_like) + len(hook_pos) + len(clips) # Verify the selected counts match the categorized lists assert details["selected_counts"]["liked_hooks"] == len(hook_like) assert details["selected_counts"]["positive_hooks"] == len(hook_pos) def test_audio_seeds_with_custom_config(self): """Test audio seeds with custom configuration.""" with patch.dict(RECS_CONFIG, { "audio_seed_boosts": { "liked_hooks": 1, "positive_hooks": 1, "liked_clips": 1, "listening_clips": 1 }, "max_audio_knn_seeds": 20 }): liked_hooks = [f"lh{i}" for i in range(10)] positive_hooks = [f"ph{i}" for i in range(10)] liked_clips = [f"lc{i}" for i in range(10)] listening_clips = [f"ls{i}" for i in range(10)] hook_like, hook_pos, clips, details = get_audio_knn_seeds( liked_hooks, positive_hooks, liked_clips, listening_clips ) total = len(hook_like) + len(hook_pos) + len(clips) assert total <= 20 # With equal weights, distribution should be more balanced counts = details["selected_counts"] values = list(counts.values()) # All counts should be within reasonable range of each other # With 20 samples across 4 categories, we expect ~5 each # but randomness can cause more variance assert max(values) - min(values) <= 6 # Allow reasonable variance if __name__ == '__main__': pytest.main([__file__, '-v'])