334 lines
9.8 KiB
Rust
334 lines
9.8 KiB
Rust
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");
|
|
}
|