path: root/packages/meshbay-node/tests/test_enrich.py

"""Tests for indexer/enrich.py — season/title corroboration and the end-to-end pool."""

import asyncio
import shutil
import subprocess
from pathlib import Path

import pytest

from meshbay_common.protocol import IndexEntry
from meshbay_node.indexer.enrich import (
    Enricher, _season_from_ancestors, _synthetic_episode_number,
    _title_from_show_siblings, _title_from_siblings,
)
from meshbay_node.media_cache import MediaCache

_HAVE_FFMPEG = shutil.which("ffmpeg") and shutil.which("ffprobe")


# ── pure helpers, no ffmpeg needed ───────────────────────────────────────────

def test_season_from_ancestors_finds_season_folder(tmp_path):
    folder = tmp_path / "Some Show" / "Season 2"
    folder.mkdir(parents=True)
    ep = folder / "01 - Episode Title.mkv"
    ep.touch()

    assert _season_from_ancestors(ep) == 2


def test_season_from_ancestors_none_when_no_season_folder(tmp_path):
    folder = tmp_path / "Movies"
    folder.mkdir()
    f = folder / "Some Movie 2015.mkv"
    f.touch()

    assert _season_from_ancestors(f) is None


def test_title_from_siblings_borrows_from_a_titled_sibling(tmp_path):
    folder = tmp_path / "Acronym Show"
    folder.mkdir()
    titled = folder / "Some.Show.Name.S01E01.720p.mkv"
    untitled = folder / "S01E02.SUBFRENCH.720p.mkv"
    titled.touch()
    untitled.touch()

    assert _title_from_siblings(untitled) == "Some Show Name"


def test_title_from_siblings_none_when_no_titled_sibling(tmp_path):
    folder = tmp_path / "Acronym Show"
    folder.mkdir()
    (folder / "S01E02.mkv").touch()

    assert _title_from_siblings(folder / "S01E02.mkv") is None


def test_title_from_siblings_ignores_a_titled_sibling_with_no_episode_number(tmp_path):
    # A Specials-shaped folder: every file parses as a "movie" (its own
    # one-off title, guessit finds no season/episode grammar at all) — none
    # of them is a representative sibling for the show's actual name.
    folder = tmp_path / "Some Show" / "Specials"
    folder.mkdir(parents=True)
    (folder / "Bonus Feature One.mkv").touch()

    assert _title_from_siblings(folder / "Bonus Feature Two.mkv") is None


def test_title_from_show_siblings_borrows_from_a_sibling_season_folder(tmp_path):
    show = tmp_path / "Some Show"
    specials = show / "Specials"
    specials.mkdir(parents=True)
    (specials / "Bonus Feature One.mkv").touch()
    season1 = show / "Season 1"
    season1.mkdir()
    (season1 / "Some Show.S01E01.mkv").touch()

    assert _title_from_show_siblings(specials / "Bonus Feature One.mkv") == "Some Show"


def test_title_from_show_siblings_none_when_no_show_folder_has_one(tmp_path):
    show = tmp_path / "Some Show"
    specials = show / "Specials"
    specials.mkdir(parents=True)
    (specials / "Bonus Feature One.mkv").touch()
    (specials / "Bonus Feature Two.mkv").touch()

    assert _title_from_show_siblings(specials / "Bonus Feature One.mkv") is None


def test_title_from_show_siblings_not_specific_to_the_word_specials(tmp_path):
    # Genericity check: the same gap (a season-like folder whose files carry
    # no episode grammar of their own) can show up under a plain numbered
    # season folder too, not only one named "Specials" — nothing in
    # _title_from_show_siblings should assume otherwise.
    show = tmp_path / "Some Show"
    season2 = show / "Season 2"
    season2.mkdir(parents=True)
    (season2 / "One-Off Episode Title.mkv").touch()
    season1 = show / "Season 1"
    season1.mkdir()
    (season1 / "Some Show.S01E01.mkv").touch()

    assert _title_from_show_siblings(season2 / "One-Off Episode Title.mkv") == "Some Show"


def test_synthetic_episode_number_is_stable_alphabetical_rank(tmp_path):
    folder = tmp_path / "Specials"
    folder.mkdir()
    (folder / "Bonus Feature One.mkv").touch()
    (folder / "Bonus Feature Three.mkv").touch()
    (folder / "Bonus Feature Two.mkv").touch()

    assert _synthetic_episode_number(folder / "Bonus Feature One.mkv") == 1
    assert _synthetic_episode_number(folder / "Bonus Feature Three.mkv") == 2
    assert _synthetic_episode_number(folder / "Bonus Feature Two.mkv") == 3


# ── end-to-end against a real (tiny, synthetic) video file ──────────────────

pytestmark_ffmpeg = pytest.mark.skipif(not _HAVE_FFMPEG, reason="ffmpeg/ffprobe not installed")


def _make_clip(path: Path) -> None:
    subprocess.run(
        ["ffmpeg", "-hide_banner", "-loglevel", "error", "-y",
         "-f", "lavfi", "-i", "testsrc=size=320x240:rate=25:duration=1",
         "-c:v", "libx264", "-preset", "ultrafast", "-an", str(path)],
        check=True, capture_output=True,
    )


@pytest.fixture
async def media_cache(tmp_path):
    c = MediaCache(db_path=tmp_path / "media_cache.db")
    await c.open()
    yield c
    await c.close()


@pytestmark_ffmpeg
@pytest.mark.asyncio
async def test_enricher_populates_fields_and_stores_thumbnail(tmp_path, media_cache):
    clip = tmp_path / "Some.Movie.2015.1080p.mkv"
    _make_clip(clip)
    entry = IndexEntry(id="fileid1", name=clip.name, path=clip.name,
                        size=clip.stat().st_size, type="video", added_at=0)

    enricher = Enricher(media_cache)
    done = asyncio.get_event_loop().create_future()

    async def on_done(file_id, fields):
        done.set_result((file_id, fields))

    enricher.spawn(entry, clip, on_done)
    file_id, fields = await asyncio.wait_for(done, timeout=30)

    assert file_id == "fileid1"
    assert fields["width"] == 320
    assert fields["height"] == 240
    assert fields["display_title"] == "Some Movie"
    assert fields.get("thumb_hash")
    stored = await media_cache.get_thumb(fields["thumb_hash"])
    assert stored is not None and len(stored) > 0


async def _run(enricher, entry, path):
    done = asyncio.get_event_loop().create_future()

    async def on_done(file_id, fields):
        done.set_result((file_id, fields))

    enricher.spawn(entry, path, on_done)
    return await asyncio.wait_for(done, timeout=30)


@pytestmark_ffmpeg
@pytest.mark.asyncio
async def test_enricher_reuses_cached_probe_and_thumbnail_without_touching_the_file_again(
        tmp_path, media_cache):
    """
    The gap this closes: duration/width/height/thumb_hash only ever lived
    on the in-memory GroupIndex entry, so every daemon restart re-ran
    ffprobe and ffmpeg over every video in every group from scratch, even
    though media_cache.db already had the answer. Proven strongly: the
    source file is deleted between the two runs, so a real second probe or
    frame grab would fail outright rather than merely being redundant.
    """
    clip = tmp_path / "Some.Movie.2015.1080p.mkv"
    _make_clip(clip)
    entry = IndexEntry(id="fileid_reuse", name=clip.name, path=clip.name,
                        size=clip.stat().st_size, type="video", added_at=0)

    first_enricher = Enricher(media_cache)
    _, first_fields = await _run(first_enricher, entry, clip)
    assert first_fields.get("thumb_hash")

    clip.unlink()
    second_enricher = Enricher(media_cache)
    _, second_fields = await _run(second_enricher, entry, clip)

    assert second_fields == first_fields


@pytestmark_ffmpeg
@pytest.mark.asyncio
async def test_a_renamed_file_re_parses_its_title_even_with_a_cached_probe(tmp_path, media_cache):
    """
    A cached probe/thumbnail must never leak into a stale title parse.
    display_title/season/episode come from guessit against the *filename*,
    which is exactly what a rename needs re-derived
    (test_rename_reenrichment.py covers the scheduling half of this) — the
    new cache-hit path added alongside them must not accidentally reuse a
    stale parse just because it took the same shortcut for duration/width/
    height. Proven the same strong way: the renamed file never exists on
    disk at all, so a correct implementation still succeeds (guessit only
    reads entry.name) while a regression that tried to re-probe or
    re-thumbnail the "new" path would fail outright.
    """
    clip = tmp_path / "Old.Name.2015.mkv"
    _make_clip(clip)
    entry = IndexEntry(id="fileid_rename", name=clip.name, path=clip.name,
                        size=clip.stat().st_size, type="video", added_at=0)

    enricher = Enricher(media_cache)
    _, first_fields = await _run(enricher, entry, clip)
    assert first_fields["display_title"] == "Old Name"

    renamed_path = tmp_path / "New.Name.2020.mkv"  # never created — proves nothing re-reads it
    renamed_entry = IndexEntry(id="fileid_rename", name=renamed_path.name,
                               path=renamed_path.name, size=entry.size, type="video", added_at=0)
    _, second_fields = await _run(enricher, renamed_entry, renamed_path)

    assert second_fields["display_title"] == "New Name"
    assert second_fields["width"] == first_fields["width"]
    assert second_fields["height"] == first_fields["height"]
    assert second_fields["thumb_hash"] == first_fields["thumb_hash"]


@pytestmark_ffmpeg
@pytest.mark.asyncio
async def test_enricher_handles_episode_with_season_from_folder(tmp_path, media_cache):
    # Filename carries only a bare episode number, no SxxExx token — guessit
    # confirmed (separately) not to find a season here at all — so the
    # season must come from the ancestor folder (§3.4's non-standard case).
    folder = tmp_path / "Some Show" / "Saison 3"
    folder.mkdir(parents=True)
    titled_sibling = folder / "Some.Show.Episode.06.mkv"
    titled_sibling.touch()
    clip = folder / "Episode.07.720p.mkv"
    _make_clip(clip)
    entry = IndexEntry(id="fileid2", name=clip.name, path=str(clip.relative_to(tmp_path)),
                        size=clip.stat().st_size, type="video", added_at=0)

    enricher = Enricher(media_cache)
    done = asyncio.get_event_loop().create_future()

    async def on_done(file_id, fields):
        done.set_result((file_id, fields))

    enricher.spawn(entry, clip, on_done)
    file_id, fields = await asyncio.wait_for(done, timeout=30)

    assert fields["display_title"] == "Some Show"
    assert fields["season"] == 3
    assert fields["episode"] == 7


@pytestmark_ffmpeg
@pytest.mark.asyncio
async def test_enricher_handles_specials_folder_with_no_episode_grammar(tmp_path, media_cache):
    # The real bug this guards: every file in a Specials folder is named
    # after its own one-off joke, not the show — guessit finds no
    # season/episode in any of them, so without the ancestor-folder check
    # each one used to fall to the movie branch and get searched against
    # TMDB as an unrelated standalone film (found live: a real show's
    # Specials folder matched several of its bonus episodes to real,
    # unrelated movies that happened to share their one-off titles).
    show = tmp_path / "Some Show"
    season1 = show / "Season 1"
    season1.mkdir(parents=True)
    (season1 / "Some Show.S01E01.mkv").touch()
    specials = show / "Specials"
    specials.mkdir()
    (specials / "Bonus Feature One.mkv").touch()
    clip = specials / "Bonus Feature Two.mkv"
    _make_clip(clip)
    entry = IndexEntry(id="fileid3", name=clip.name, path=str(clip.relative_to(tmp_path)),
                        size=clip.stat().st_size, type="video", added_at=0)

    enricher = Enricher(media_cache)
    _, fields = await _run(enricher, entry, clip)

    # Not "Bonus Feature Two" — that's this one Special's own title, and
    # matching it against TMDB by itself is exactly the bug. The show's
    # name, from its ordinary season folder next door.
    assert fields["display_title"] == "Some Show"
    assert fields["season"] == 0
    assert fields["episode"] is not None