Added the playlist generator

shanty-playlist/tests/unit.rs

@@ -0,0 +1,333 @@
+use std::collections::{HashMap, HashSet};
+
+use chrono::Utc;
+use shanty_data::PopularTrack;
+use shanty_db::entities::track::Model as Track;
+use shanty_playlist::ordering::{interleave_artists, shuffle};
+use shanty_playlist::scoring::score_tracks;
+use shanty_playlist::selection::generate_playlist;
+use shanty_playlist::types::Candidate;
+
+fn make_track(id: i32, title: &str, artist_id: Option<i32>, mbid: Option<&str>) -> Track {
+    let now = Utc::now().naive_utc();
+    Track {
+        id,
+        file_path: format!("/music/{title}.opus"),
+        title: Some(title.to_string()),
+        artist: Some("Test Artist".to_string()),
+        album: Some("Test Album".to_string()),
+        album_artist: None,
+        track_number: Some(id),
+        disc_number: None,
+        duration: Some(200.0),
+        codec: None,
+        bitrate: None,
+        genre: None,
+        year: None,
+        musicbrainz_id: mbid.map(String::from),
+        file_size: 1000,
+        fingerprint: None,
+        artist_id,
+        album_id: None,
+        added_at: now,
+        updated_at: now,
+        file_mtime: None,
+    }
+}
+
+fn make_candidate(id: i32, artist: &str, score: f64) -> Candidate {
+    Candidate {
+        score,
+        artist: artist.to_string(),
+        artist_mbid: Some(format!("mbid-{artist}")),
+        track_id: id,
+        file_path: format!("/music/{id}.opus"),
+        title: Some(format!("Track {id}")),
+        album: None,
+        duration: Some(200.0),
+    }
+}
+
+// --- Scoring tests ---
+
+#[test]
+fn test_score_tracks_basic() {
+    let artists = vec![("artist-1".to_string(), "Artist One".to_string(), 1.0)];
+
+    let tracks: Vec<Track> = (1..=5)
+        .map(|i| make_track(i, &format!("Song {i}"), Some(1), None))
+        .collect();
+
+    let top_tracks = vec![
+        PopularTrack {
+            name: "Song 1".to_string(),
+            mbid: None,
+            playcount: 1000,
+        },
+        PopularTrack {
+            name: "Song 2".to_string(),
+            mbid: None,
+            playcount: 500,
+        },
+        PopularTrack {
+            name: "Song 3".to_string(),
+            mbid: None,
+            playcount: 100,
+        },
+    ];
+
+    let mut tracks_map = HashMap::new();
+    tracks_map.insert("artist-1".to_string(), tracks);
+
+    let mut top_map = HashMap::new();
+    top_map.insert("artist-1".to_string(), top_tracks);
+
+    let scored = score_tracks(&artists, &tracks_map, &top_map, 5);
+
+    // Should have 3 tracks (only ones matching top tracks)
+    assert_eq!(scored.len(), 3);
+
+    // Higher playcount should yield higher score
+    let scores: Vec<f64> = scored.iter().map(|t| t.score).collect();
+    let max_score = scores.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
+    assert!(max_score > 0.0);
+}
+
+#[test]
+fn test_score_tracks_no_top_tracks_uses_uniform() {
+    let artists = vec![("artist-1".to_string(), "Artist One".to_string(), 0.8)];
+
+    let tracks: Vec<Track> = (1..=3)
+        .map(|i| make_track(i, &format!("Song {i}"), Some(1), None))
+        .collect();
+
+    let mut tracks_map = HashMap::new();
+    tracks_map.insert("artist-1".to_string(), tracks);
+
+    let top_map = HashMap::new(); // No top tracks
+
+    let scored = score_tracks(&artists, &tracks_map, &top_map, 5);
+
+    // All 3 tracks should be included with uniform scoring
+    assert_eq!(scored.len(), 3);
+    // All should have the same score (similarity only)
+    let first_score = scored[0].score;
+    for t in &scored {
+        assert!((t.score - first_score).abs() < 1e-10);
+    }
+}
+
+#[test]
+fn test_score_tracks_per_artist_cap() {
+    let artists = vec![("artist-1".to_string(), "Artist One".to_string(), 1.0)];
+
+    // 50 tracks
+    let tracks: Vec<Track> = (1..=50)
+        .map(|i| make_track(i, &format!("Song {i}"), Some(1), None))
+        .collect();
+
+    let top_tracks: Vec<PopularTrack> = (1..=50)
+        .map(|i| PopularTrack {
+            name: format!("Song {i}"),
+            mbid: None,
+            playcount: (50 - i + 1) as u64 * 100,
+        })
+        .collect();
+
+    let mut tracks_map = HashMap::new();
+    tracks_map.insert("artist-1".to_string(), tracks);
+    let mut top_map = HashMap::new();
+    top_map.insert("artist-1".to_string(), top_tracks);
+
+    // bias 10 → cap = 10
+    let scored = score_tracks(&artists, &tracks_map, &top_map, 10);
+    assert!(scored.len() <= 10);
+
+    // bias 0 → no cap
+    let scored_no_cap = score_tracks(&artists, &tracks_map, &top_map, 0);
+    assert_eq!(scored_no_cap.len(), 50);
+}
+
+#[test]
+fn test_similarity_transform() {
+    // Higher match_score should produce higher similarity
+    let artists = vec![
+        ("high".to_string(), "High".to_string(), 0.9),
+        ("low".to_string(), "Low".to_string(), 0.1),
+    ];
+
+    let track_high = make_track(1, "Song", Some(1), None);
+    let track_low = make_track(2, "Song", Some(2), None);
+
+    let mut tracks_map = HashMap::new();
+    tracks_map.insert("high".to_string(), vec![track_high]);
+    tracks_map.insert("low".to_string(), vec![track_low]);
+
+    let scored = score_tracks(&artists, &tracks_map, &HashMap::new(), 5);
+    assert_eq!(scored.len(), 2);
+
+    let high_score = scored
+        .iter()
+        .find(|t| t.artist == "High")
+        .unwrap()
+        .similarity;
+    let low_score = scored
+        .iter()
+        .find(|t| t.artist == "Low")
+        .unwrap()
+        .similarity;
+    assert!(high_score > low_score);
+}
+
+// --- Selection tests ---
+
+#[test]
+fn test_generate_playlist_basic() {
+    let candidates: Vec<Candidate> = (1..=20)
+        .map(|i| make_candidate(i, &format!("Artist{}", i % 4), 1.0))
+        .collect();
+
+    let seeds = HashSet::new();
+    let result = generate_playlist(&candidates, 10, &seeds);
+
+    assert_eq!(result.len(), 10);
+}
+
+#[test]
+fn test_generate_playlist_respects_count() {
+    let candidates: Vec<Candidate> = (1..=5).map(|i| make_candidate(i, "Artist", 1.0)).collect();
+
+    let seeds = HashSet::new();
+    let result = generate_playlist(&candidates, 3, &seeds);
+    assert_eq!(result.len(), 3);
+}
+
+#[test]
+fn test_generate_playlist_not_more_than_available() {
+    let candidates: Vec<Candidate> = (1..=3).map(|i| make_candidate(i, "Artist", 1.0)).collect();
+
+    let seeds = HashSet::new();
+    let result = generate_playlist(&candidates, 100, &seeds);
+    assert_eq!(result.len(), 3);
+}
+
+#[test]
+fn test_generate_playlist_empty_candidates() {
+    let candidates: Vec<Candidate> = vec![];
+    let seeds = HashSet::new();
+    let result = generate_playlist(&candidates, 10, &seeds);
+    assert!(result.is_empty());
+}
+
+#[test]
+fn test_generate_playlist_per_artist_cap() {
+    // 20 tracks from one artist, 5 from another
+    let mut candidates: Vec<Candidate> = (1..=20)
+        .map(|i| make_candidate(i, "Prolific", 1.0))
+        .collect();
+    candidates.extend((21..=25).map(|i| make_candidate(i, "Minor", 1.0)));
+
+    let seeds = HashSet::new();
+    let result = generate_playlist(&candidates, 15, &seeds);
+
+    let prolific_count = result.iter().filter(|c| c.artist == "Prolific").count();
+    let minor_count = result.iter().filter(|c| c.artist == "Minor").count();
+
+    // With 2 artists and n=15, cap = ceil(15/2) = 8.
+    // Minor only has 5 tracks, so Prolific fills the rest via fallback.
+    // Key check: both artists are represented.
+    assert!(
+        minor_count >= 1,
+        "Minor should get at least 1 track, got {minor_count}"
+    );
+    assert!(
+        prolific_count >= 1,
+        "Prolific should get at least 1 track, got {prolific_count}"
+    );
+    assert_eq!(result.len(), 15);
+}
+
+#[test]
+fn test_generate_playlist_seed_enforcement() {
+    // Many tracks from "Other" with high scores, few from "Seed" with low scores
+    let mut candidates: Vec<Candidate> =
+        (1..=50).map(|i| make_candidate(i, "Other", 10.0)).collect();
+    candidates.push(make_candidate(51, "Seed", 0.01));
+    candidates.push(make_candidate(52, "Seed", 0.01));
+
+    let mut seeds = HashSet::new();
+    seeds.insert("Seed".to_string());
+
+    let result = generate_playlist(&candidates, 10, &seeds);
+    let seed_count = result.iter().filter(|c| c.artist == "Seed").count();
+
+    // seed_min = (10/10).max(1) = 1, so at least 1 seed track
+    assert!(
+        seed_count >= 1,
+        "Expected at least 1 seed track, got {seed_count}"
+    );
+}
+
+// --- Ordering tests ---
+
+#[test]
+fn test_interleave_no_back_to_back() {
+    let tracks: Vec<Candidate> = vec![
+        make_candidate(1, "A", 1.0),
+        make_candidate(2, "A", 1.0),
+        make_candidate(3, "A", 1.0),
+        make_candidate(4, "B", 1.0),
+        make_candidate(5, "B", 1.0),
+        make_candidate(6, "B", 1.0),
+    ];
+
+    let result = interleave_artists(tracks);
+    assert_eq!(result.len(), 6);
+
+    // Check no back-to-back same artist
+    for window in result.windows(2) {
+        assert_ne!(
+            window[0].artist, window[1].artist,
+            "Back-to-back: {} at positions",
+            window[0].artist
+        );
+    }
+}
+
+#[test]
+fn test_interleave_single_artist() {
+    let tracks: Vec<Candidate> = (1..=5).map(|i| make_candidate(i, "Solo", 1.0)).collect();
+
+    let result = interleave_artists(tracks);
+    assert_eq!(result.len(), 5);
+    // All same artist, so back-to-back is unavoidable — just check count
+}
+
+#[test]
+fn test_shuffle_preserves_count() {
+    let tracks: Vec<Candidate> = (1..=10).map(|i| make_candidate(i, "Artist", 1.0)).collect();
+
+    let result = shuffle(tracks);
+    assert_eq!(result.len(), 10);
+}
+
+#[test]
+fn test_interleave_many_artists() {
+    let mut tracks = Vec::new();
+    for artist_idx in 0..5 {
+        for track_idx in 0..4 {
+            let id = artist_idx * 4 + track_idx + 1;
+            tracks.push(make_candidate(id, &format!("Artist{artist_idx}"), 1.0));
+        }
+    }
+
+    let result = interleave_artists(tracks);
+    assert_eq!(result.len(), 20);
+
+    // Count back-to-back violations (should be 0 with 5 artists)
+    let violations = result
+        .windows(2)
+        .filter(|w| w[0].artist == w[1].artist)
+        .count();
+    assert_eq!(violations, 0, "Expected no back-to-back with 5 artists");
+}