From c16e9f1035efed36796ef9f3b84976c7228aab3a Mon Sep 17 00:00:00 2001 From: Jonathan Date: Fri, 10 Jul 2026 23:17:06 +0200 Subject: [PATCH] feat: MusicBrainz-aware intake and completeness check against canonical track count --- worker/lyra_worker/pipeline.py | 10 ++++- worker/lyra_worker/resolver.py | 9 ++++ worker/lyra_worker/types.py | 1 + worker/tests/test_intake.py | 78 ++++++++++++++++++++++++++++++++++ 4 files changed, 96 insertions(+), 2 deletions(-) create mode 100644 worker/lyra_worker/resolver.py create mode 100644 worker/tests/test_intake.py diff --git a/worker/lyra_worker/pipeline.py b/worker/lyra_worker/pipeline.py index adeb2fa..bffb25d 100644 --- a/worker/lyra_worker/pipeline.py +++ b/worker/lyra_worker/pipeline.py @@ -101,6 +101,7 @@ def run_pipeline( conn: psycopg.Connection, job_id: str, adapters: Sequence[SourceAdapter], + resolver=None, min_confidence: float = 0.7, dest_root: str = "/music", ) -> None: @@ -112,6 +113,10 @@ def run_pipeline( # 1. intake _set_state(conn, job_id, "matching", "intake") target = _load_target(conn, job_id) + if resolver is not None: + resolved = resolver.resolve(target.artist, target.album) + if resolved is not None: + target = resolved if _already_in_library(conn, target): _set_state(conn, job_id, "imported", "import") with conn.cursor() as cur: @@ -161,9 +166,10 @@ def run_pipeline( _fail(conn, job_id, "all downloads failed") return - # 5. tag (integrity check; real tagging is Plan 4) + # 5. tag (integrity check against the canonical track count when known) _set_state(conn, job_id, "tagging", "tag") - if result.track_count != winner.track_count: + expected = target.track_count if target.track_count is not None else winner.track_count + if result.track_count < expected: _fail(conn, job_id, "incomplete download") return diff --git a/worker/lyra_worker/resolver.py b/worker/lyra_worker/resolver.py new file mode 100644 index 0000000..c46d7e1 --- /dev/null +++ b/worker/lyra_worker/resolver.py @@ -0,0 +1,9 @@ +from typing import Protocol + +from lyra_worker.types import MBTarget + + +class MbResolver(Protocol): + def resolve(self, artist: str, album: str) -> MBTarget | None: + """Return a canonical MBTarget (with track_count) or None if no confident match.""" + ... diff --git a/worker/lyra_worker/types.py b/worker/lyra_worker/types.py index a67708a..928a902 100644 --- a/worker/lyra_worker/types.py +++ b/worker/lyra_worker/types.py @@ -7,6 +7,7 @@ class MBTarget: album: str track_count: int | None = None total_duration_s: int | None = None + year: int | None = None @dataclass(frozen=True) diff --git a/worker/tests/test_intake.py b/worker/tests/test_intake.py new file mode 100644 index 0000000..5bc4583 --- /dev/null +++ b/worker/tests/test_intake.py @@ -0,0 +1,78 @@ +from lyra_worker.adapters.youtube import YouTubeAdapter +from lyra_worker.claim import claim_next +from lyra_worker.pipeline import run_pipeline +from lyra_worker.types import MBTarget +from tests.conftest import insert_request +from tests.test_youtube_adapter import FakeYtClient + + +class FakeResolver: + def __init__(self, target): + self._target = target + + def resolve(self, artist, album): + return self._target + + +def _job_state(conn, job_id): + with conn.cursor() as cur: + cur.execute('SELECT state FROM "Job" WHERE id = %s', (job_id,)) + return cur.fetchone()[0] + + +def test_single_track_download_rejected_against_musicbrainz_count(conn): + # MusicBrainz says Continuum has 12 tracks; a 1-track YouTube "full album" video must be rejected. + resolver = FakeResolver(MBTarget(artist="John Mayer", album="Continuum", track_count=12)) + yt = [{"source_ref": "yt1", "title": "Continuum", "artist": "John Mayer", "track_count": 1}] + # FakeYtClient.download reports the same 1-track count for the single-video candidate: + client = FakeYtClient(results=yt) + client.download = lambda source_ref, dest, on_progress: (on_progress(1.0), {"track_count": 1, "path": dest})[1] + + job_id = insert_request(conn, artist="john mayer", album="continuum") + claim_next(conn) + run_pipeline(conn, job_id, [YouTubeAdapter(client)], resolver=resolver, dest_root="/tmp/lib") + + assert _job_state(conn, job_id) == "needs_attention" + + +def test_complete_download_imports_with_musicbrainz_count(conn): + resolver = FakeResolver(MBTarget(artist="John Mayer", album="Continuum", track_count=12)) + yt = [{"source_ref": "yt1", "title": "Continuum", "artist": "John Mayer", "track_count": 12}] + client = FakeYtClient(results=yt) + client.download = lambda source_ref, dest, on_progress: (on_progress(1.0), {"track_count": 12, "path": dest})[1] + + job_id = insert_request(conn, artist="john mayer", album="continuum") + claim_next(conn) + run_pipeline(conn, job_id, [YouTubeAdapter(client)], resolver=resolver, dest_root="/tmp/lib") + + assert _job_state(conn, job_id) == "imported" + with conn.cursor() as cur: + # canonical artist/album from MusicBrainz are what get stored + cur.execute('SELECT artist, album FROM "LibraryItem"') + assert cur.fetchone() == ("John Mayer", "Continuum") + + +def test_no_resolver_preserves_legacy_behavior(conn): + # Without a resolver, target.track_count stays None and the candidate's own count is used. + yt = [{"source_ref": "yt1", "title": "Continuum", "artist": "John Mayer", "track_count": 1}] + client = FakeYtClient(results=yt) + client.download = lambda source_ref, dest, on_progress: (on_progress(1.0), {"track_count": 1, "path": dest})[1] + + job_id = insert_request(conn, artist="John Mayer", album="Continuum") + claim_next(conn) + run_pipeline(conn, job_id, [YouTubeAdapter(client)], dest_root="/tmp/lib") # no resolver + + assert _job_state(conn, job_id) == "imported" # 1 == 1, passes as before + + +def test_resolver_returning_none_falls_back_to_request(conn): + resolver = FakeResolver(None) + yt = [{"source_ref": "yt1", "title": "Continuum", "artist": "John Mayer", "track_count": 1}] + client = FakeYtClient(results=yt) + client.download = lambda source_ref, dest, on_progress: (on_progress(1.0), {"track_count": 1, "path": dest})[1] + + job_id = insert_request(conn, artist="John Mayer", album="Continuum") + claim_next(conn) + run_pipeline(conn, job_id, [YouTubeAdapter(client)], resolver=resolver, dest_root="/tmp/lib") + + assert _job_state(conn, job_id) == "imported" # falls back to raw target, 1==1