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

"""SQL рекомендателя на настоящей схеме (CI-job `ci / db`; схему заранее накатывают миграции api).

Локально: TEST_DATABASE_URL=postgresql://… uv run pytest tests/test_db_integration.py
"""

import os
import time
import uuid
from typing import Any

import numpy as np
import psycopg
import pytest

from pipu_rec.analysis.features import Features
from pipu_rec.analysis.worker import Worker
from pipu_rec.train.pipeline import load_interactions, train_all
from pipu_rec.wave.engine import Profile, Session, build_batch
from pipu_rec.wave.store import load_index, load_profile

DSN = os.environ.get("TEST_DATABASE_URL", "")
pytestmark = pytest.mark.skipif(not DSN, reason="TEST_DATABASE_URL is not set")


def uid() -> str:
    return str(uuid.uuid4())


@pytest.fixture()
def world() -> dict[str, Any]:
    """Хранилище, 2 исполнителя × 10 треков, 6 слушателей с событиями и реакциями."""
    with psycopg.connect(DSN, autocommit=True) as conn:
        src = uid()
        conn.execute(
            "INSERT INTO library_sources (id, name, type, local_path, status) VALUES (%s, 'it', 'local', '', 'mounted')",
            (src,),
        )
        tracks: list[str] = []
        for a in range(2):
            artist = uid()
            conn.execute(
                "INSERT INTO artists (id, name, sort_name, name_norm) VALUES (%s, %s, %s, %s)",
                (artist, f"A{a}", f"A{a}", f"a{a}-{artist[:6]}"),
            )
            album = uid()
            conn.execute(
                "INSERT INTO albums (id, title, title_norm, album_artist_id) VALUES (%s, %s, %s, %s)",
                (album, f"Al{a}", f"al{a}", artist),
            )
            for i in range(10):
                t, f = uid(), uid()
                conn.execute(
                    "INSERT INTO tracks (id, title, title_norm, album_id, duration_ms) VALUES (%s, %s, %s, %s, 200000)",
                    (t, f"T{a}{i}", f"t{a}{i}", album),
                )
                conn.execute(
                    "INSERT INTO media_files (id, source_id, rel_path, size_bytes, mtime, format, duration_ms, track_id) "
                    "VALUES (%s, %s, %s, 1, now(), 'flac', 200000, %s)",
                    (f, src, f"{a}/{i}.flac", t),
                )
                conn.execute("UPDATE tracks SET primary_file_id = %s WHERE id = %s", (f, t))
                conn.execute(
                    "INSERT INTO track_artists (track_id, artist_id, role, position) VALUES (%s, %s, 'main', 0)",
                    (t, artist),
                )
                tracks.append(t)
        users = []
        for u in range(6):
            user = uid()
            conn.execute(
                "INSERT INTO users (id, username, password_hash) VALUES (%s, %s, 'x')",
                (user, f"u{u}-{user[:6]}"),
            )
            users.append(user)
            group = tracks[:10] if u % 2 == 0 else tracks[10:]
            for k, t in enumerate(group):
                kind = "play_end" if k < 8 else "skip"
                conn.execute(
                    "INSERT INTO play_events (id, user_id, track_id, type, occurred_at, played_ms, client_type) "
                    "VALUES (%s, %s, %s, %s, now() - interval '20 days', %s, 'web')",
                    (uid(), user, t, kind, 200000 if kind == "play_end" else 3000),
                )
                conn.execute(
                    "INSERT INTO user_track_stats (user_id, track_id, play_count, complete_count, last_played_at) "
                    "VALUES (%s, %s, 1, %s, now() - interval '20 days')",
                    (user, t, 1 if kind == "play_end" else 0),
                )
            conn.execute(
                "INSERT INTO track_reactions (user_id, track_id, reaction) VALUES (%s, %s, 'like')",
                (user, group[0]),
            )
        return {"tracks": tracks, "users": users}


def test_analysis_worker_saves_features(world: dict[str, Any]) -> None:
    w = Worker(DSN, "/library", "/library-host", 1)
    emb = np.zeros(256, dtype=np.float32)
    emb[0] = 1.0
    feats = Features(
        bpm=120.0,
        bpm_confidence=0.9,
        key="C",
        scale="major",
        loudness_lufs=-12.0,
        energy=0.6,
        danceability=0.5,
        valence=0.5,
        acousticness=0.2,
        intro_ms=1000,
        embedding=emb,
    )
    with psycopg.connect(DSN, autocommit=True) as conn:
        jobs = w.claim(conn, 5)
        assert jobs, "every new track is queued for analysis by the trigger"
        w.save(conn, jobs[0].track_id, feats)
        w.fail(conn, jobs[1].track_id, "boom")
        w.requeue_stale(conn)
        row = conn.execute(
            "SELECT bpm FROM track_features WHERE track_id = %s", (jobs[0].track_id,)
        ).fetchone()
        assert row is not None and row[0] == pytest.approx(120.0)


def test_training_and_wave_queries(world: dict[str, Any]) -> None:
    with psycopg.connect(DSN) as conn:
        assert len(load_interactions(conn)) >= 60
    result = train_all(DSN, now=time.time())
    assert result.get("activated") is True, result
    with psycopg.connect(DSN) as conn:
        index = load_index(conn)
        assert set(world["tracks"]) <= set(index.ids)
        assert index.has_cf.sum() >= 20
        profile: Profile = load_profile(conn, world["users"][0])
        assert profile.taste_cf is not None
        assert profile.recent_likes
        active = conn.execute(
            "SELECT count(*) FROM model_versions WHERE kind = 'als' AND is_active"
        ).fetchone()
        assert active is not None and active[0] == 1
    picks = build_batch(index, profile, Session(), [], n=5, rng=np.random.default_rng(0))
    assert len(picks) == 5