pipu-music/services/recommender/tests/test_engine.py

import numpy as np
import pytest

from pipu_rec.wave.contexts import load_contexts
from pipu_rec.wave.engine import CatalogIndex, Profile, Session, build_batch, taste_centroids, update_session


def library(n_per_cluster: int = 20, seed: int = 0) -> CatalogIndex:
    """Три «жанра» — кластера эмбеддингов; в кластере 5 исполнителей по 4 трека."""
    rng = np.random.default_rng(seed)
    centers = rng.normal(size=(3, 32))
    ids, artists, embs = [], [], []
    feats: dict[str, list[float | None]] = {
        "bpm": [],
        "energy": [],
        "danceability": [],
        "valence": [],
        "acousticness": [],
    }
    for c in range(3):
        for k in range(n_per_cluster):
            ids.append(f"t{c}-{k}")
            artists.append(f"a{c}-{k // 4}")
            embs.append((centers[c] + 0.3 * rng.normal(size=32)).astype(np.float32))
            feats["bpm"].append([90.0, 128.0, 172.0][c])
            feats["energy"].append([0.2, 0.5, 0.85][c])
            feats["danceability"].append(0.7)
            feats["valence"].append(0.5)
            feats["acousticness"].append(None)
    return CatalogIndex.build(ids, artists, embs, feats, plays=[1] * len(ids), dim=32)


def profile(**kw: object) -> Profile:
    base = dict(
        scores={},
        disliked=set(),
        disliked_artists=set(),
        recent=set(),
        played=set(),
        liked_recent=set(),
        rediscover=set(),
        exploration=0.0,
    )
    base.update(kw)
    return Profile(**base)  # type: ignore[arg-type]


def cluster(track_id: str) -> int:
    return int(track_id[1])


def test_cold_start_returns_a_diverse_batch() -> None:
    idx = library()
    picks = build_batch(idx, profile(), Session(), [], n=5, rng=np.random.default_rng(1))
    assert len(picks) == 5
    assert len({p.id for p in picks}) == 5
    assert len({idx.artist[idx.pos[p.id]] for p in picks}) == 5, "one track per artist in a batch"


def test_taste_pulls_towards_liked_cluster() -> None:
    idx = library()
    liked = {f"t2-{k}": 3.0 for k in range(16)}
    picks = build_batch(
        idx, profile(scores=liked, played=set(liked)), Session(), [], n=5, rng=np.random.default_rng(2)
    )
    assert sum(cluster(p.id) == 2 for p in picks) >= 4


def test_dislikes_recent_and_served_are_excluded() -> None:
    idx = library(5)
    p = profile(disliked={"t0-0"}, disliked_artists={"a1-0"}, recent={"t2-0"})
    s = Session(served=["t2-1"])
    for seed in range(5):
        ids = {x.id for x in build_batch(idx, p, s, [], n=5, rng=np.random.default_rng(seed))}
        assert not ids & {"t0-0", "t2-0", "t2-1"}
        assert not any(idx.artist[idx.pos[i]] == "a1-0" for i in ids), "disliked artist"


def test_running_context_prefers_fast_energetic_tracks() -> None:
    idx = library()
    running = load_contexts().activities["running"].rules
    picks = build_batch(idx, profile(), Session(), running, n=5, rng=np.random.default_rng(3))
    assert all(cluster(p.id) == 2 for p in picks), [p.id for p in picks]


def test_half_time_counts_for_running() -> None:
    from pipu_rec.wave.contexts import fit

    running = load_contexts().activities["running"].rules
    f = fit(
        {
            "bpm": np.array([85.0, 120.0]),
            "energy": np.array([0.8, 0.8]),
            "danceability": np.array([0.7, 0.7]),
            "valence": np.array([np.nan, np.nan]),
            "acousticness": np.array([np.nan, np.nan]),
        },
        running,
    )
    assert f[0] > 0.9 > f[1]


def test_session_vector_follows_reactions_and_changes_course() -> None:
    idx = library()
    liked = {f"t0-{k}": 3.0 for k in range(10)} | {f"t2-{k}": 3.0 for k in range(10)}
    cents = taste_centroids(idx, profile(scores=liked))
    s = Session()
    update_session(idx, s, [("t1-0", "like"), ("t1-1", "complete")], cents)
    assert s.vec is not None
    without = build_batch(idx, profile(scores=liked), Session(), [], n=5, rng=np.random.default_rng(4))
    picks = build_batch(idx, profile(scores=liked), s, [], n=5, rng=np.random.default_rng(4))
    assert sum(cluster(p.id) == 1 for p in picks) > sum(cluster(p.id) == 1 for p in without)
    assert sum(cluster(p.id) == 1 for p in picks) >= 2, "the session pulls the wave towards its reactions"

    update_session(idx, s, [("t1-2", "skip_early"), ("t1-3", "skip_early"), ("t1-4", "dislike")], cents)
    assert s.skips_in_row == 0, "course changed after three skips"
    vec = np.array(s.vec)
    assert float(vec @ idx.emb[idx.pos["t1-0"]]) < 0.5


def test_exploration_brings_unheard_tracks() -> None:
    idx = library()
    heard = {f"t2-{k}" for k in range(20)}
    liked = {t: 3.0 for t in heard}
    picks = build_batch(
        idx,
        profile(scores=liked, played=heard, exploration=0.4),
        Session(),
        [],
        n=5,
        rng=np.random.default_rng(5),
    )
    assert sum(p.source == "fresh" for p in picks) >= 2


def test_bad_context_file_is_rejected(tmp_path) -> None:  # type: ignore[no-untyped-def]
    bad = tmp_path / "c.yaml"
    bad.write_text(
        "activities: {x: {icon: a, rules: [{feature: tempo, min: 1, sigma: 1}]}}\nmoods: {}\n", "utf-8"
    )
    with pytest.raises(ValueError):
        load_contexts(str(bad))