diff --git a/comet/services/filtering.py b/comet/services/filtering.py index b844e5e9..584cae64 100644 --- a/comet/services/filtering.py +++ b/comet/services/filtering.py @@ -73,7 +73,7 @@ def alternate_title_match(torrent_title: str, title: str, aliases) -> bool: continue try: - parsed_segment = _parse_with_cache(segment) + parsed_segment = parse_with_cache(segment) except ValidationError: continue @@ -208,7 +208,7 @@ def _clone_parsed(parsed): return clone -def _parse_with_cache(title: str): +def parse_with_cache(title: str): if _PARSE_CACHE_SIZE <= 0 or _PARSE_CACHE_EFFECTIVE_SHARDS <= 0: return parse(title) @@ -342,7 +342,7 @@ def filter_worker( # temp fix while waiting for RTN to fix their parsing try: - parsed = _parse_with_cache(torrent_title) + parsed = parse_with_cache(torrent_title) except ValidationError: _log_exclusion(f"❌ Rejected (Parse Error) | {torrent_title}") continue diff --git a/comet/services/orchestration.py b/comet/services/orchestration.py index aa436465..730da4e0 100644 --- a/comet/services/orchestration.py +++ b/comet/services/orchestration.py @@ -1,6 +1,7 @@ import asyncio import time +from pydantic import ValidationError from RTN import DefaultRanking, ParsedData from comet.core.execution import get_executor @@ -9,7 +10,7 @@ from comet.core.scrape import ScrapeContext from comet.scrapers.manager import scraper_manager from comet.scrapers.models import ScrapeRequest -from comet.services.filtering import filter_worker +from comet.services.filtering import TitleMatcher, filter_worker, parse_with_cache from comet.services.ranking import rank_worker from comet.services.torrent_manager import torrent_update_queue from comet.utils.languages import select_indexer_titles @@ -202,10 +203,10 @@ async def get_cached_torrents(self): for cache_media_id in self.cache_media_ids ) ) - for cache_media_id, cache_rows in zip(self.cache_media_ids, cache_row_groups): - if cache_rows and cache_media_id == self.media_only_id: - self.primary_cached = True - rows.extend(cache_rows) + for cache_media_id, cache_rows in zip( + self.cache_media_ids, cache_row_groups, strict=True + ): + rows.extend((cache_media_id, row) for row in cache_rows) if rows: best_rows = {} @@ -229,15 +230,23 @@ def row_priority(row): updated_at, ) - for row in rows: + for cache_media_id, row in rows: info_hash = row["info_hash"] current = best_rows.get(info_hash) - if current is None or row_priority(row) > row_priority(current): - best_rows[info_hash] = row + if current is None or row_priority(row) > row_priority(current[1]): + best_rows[info_hash] = (cache_media_id, row) rows = list(best_rows.values()) - for row in rows: + title_matcher = TitleMatcher( + self.title, + self.year, + self.year_end, + self.media_type, + self.aliases, + ) + + for cache_media_id, row in rows: parsed_data = load_cached_parsed(row["parsed_json"]) if parsed_data is None: logger.warning( @@ -246,6 +255,22 @@ def row_priority(row): continue ensure_multi_language(parsed_data) + torrent_title = row["title"] + if not isinstance(torrent_title, str) or not torrent_title: + continue + try: + parsed_title = parse_with_cache(torrent_title) + except ValidationError: + logger.warning( + f"Skipping torrent cache row with invalid title: {row['info_hash']}" + ) + continue + + if not parsed_title.parsed_title or not title_matcher.matches( + torrent_title, parsed_title.parsed_title, parsed_title.year + ): + continue + target_season = self.search_season if ( target_season is not None @@ -277,6 +302,8 @@ def row_priority(row): "parsed": parsed_data, "updatedAt": row["updated_at"], } + if cache_media_id == self.media_only_id: + self.primary_cached = True def _append_cache_file_infos(self, file_infos: list[dict], torrent: dict): parsed = torrent["parsed"] diff --git a/tests/test_cached_title_revalidation.py b/tests/test_cached_title_revalidation.py new file mode 100644 index 00000000..791c83d7 --- /dev/null +++ b/tests/test_cached_title_revalidation.py @@ -0,0 +1,144 @@ +import unittest +from unittest.mock import patch + +from RTN import parse + +from comet.services.orchestration import TorrentManager + + +class CachedTitleRevalidationTests(unittest.IsolatedAsyncioTestCase): + @staticmethod + def _row(title: str, info_hash: str) -> dict: + return { + "info_hash": info_hash, + "file_index": 0, + "title": title, + "seeders": 1, + "size": 1_000, + "tracker": "cache", + "sources_json": "[]", + "parsed_json": parse(title).model_dump_json(), + "episode": None, + "updated_at": 1, + } + + @staticmethod + def _manager() -> TorrentManager: + return TorrentManager( + media_type="movie", + media_full_id="tt2250912", + media_only_id="tt2250912", + title="Spider-Man: Homecoming", + year=2017, + year_end=None, + season=None, + episode=None, + aliases={}, + remove_adult_content=False, + ) + + async def test_mismatched_cached_title_is_rejected(self): + manager = self._manager() + wrong_hash = "a" * 40 + wrong = self._row( + "Spider-Man.Into.the.Spider-Verse.2018.2160p.REMUX.HEVC.DV.mkv", + wrong_hash, + ) + + with patch.object(manager, "_fetch_cached_rows", return_value=[wrong]): + await manager.get_cached_torrents() + + self.assertNotIn(wrong_hash, manager.torrents) + self.assertFalse(manager.primary_cached) + + async def test_cached_row_with_wrong_year_is_rejected(self): + manager = self._manager() + wrong_year_hash = "e" * 40 + wrong_year = self._row( + "Spider-Man.Homecoming.2019.2160p.BluRay.REMUX.HEVC.mkv", + wrong_year_hash, + ) + + with patch.object(manager, "_fetch_cached_rows", return_value=[wrong_year]): + await manager.get_cached_torrents() + + self.assertNotIn(wrong_year_hash, manager.torrents) + self.assertFalse(manager.primary_cached) + + async def test_cached_row_revalidates_legacy_parsed_data_from_raw_title(self): + manager = self._manager() + missing_title_hash = "c" * 40 + row = self._row( + "Spider-Man.Homecoming.2017.2160p.BluRay.REMUX.HEVC.mkv", + missing_title_hash, + ) + row["parsed_json"] = '{"raw_title":"Spider-Man.Homecoming.2017.mkv"}' + + with patch.object(manager, "_fetch_cached_rows", return_value=[row]): + await manager.get_cached_torrents() + + self.assertIn(missing_title_hash, manager.torrents) + self.assertTrue(manager.primary_cached) + + async def test_cached_row_does_not_trust_persisted_parsed_title(self): + manager = self._manager() + inconsistent_hash = "d" * 40 + row = self._row( + "Spider-Man.Into.the.Spider-Verse.2018.2160p.REMUX.HEVC.DV.mkv", + inconsistent_hash, + ) + row["parsed_json"] = parse( + "Spider-Man.Homecoming.2017.2160p.BluRay.REMUX.HEVC.mkv" + ).model_dump_json() + + with patch.object(manager, "_fetch_cached_rows", return_value=[row]): + await manager.get_cached_torrents() + + self.assertNotIn(inconsistent_hash, manager.torrents) + self.assertFalse(manager.primary_cached) + + async def test_matching_cached_title_still_counts_as_primary_cache(self): + manager = self._manager() + right_hash = "b" * 40 + right = self._row( + "Spider-Man.Homecoming.2017.2160p.BluRay.REMUX.HEVC.mkv", + right_hash, + ) + + with patch.object(manager, "_fetch_cached_rows", return_value=[right]): + await manager.get_cached_torrents() + + self.assertIn(right_hash, manager.torrents) + self.assertTrue(manager.primary_cached) + + async def test_secondary_duplicate_does_not_count_as_primary_cache(self): + manager = self._manager() + manager.cache_media_ids = [manager.media_only_id, "kitsu:456"] + duplicate_hash = "f" * 40 + primary = self._row( + "Spider-Man.Into.the.Spider-Verse.2018.2160p.REMUX.HEVC.DV.mkv", + duplicate_hash, + ) + primary["updated_at"] = 1 + secondary = self._row( + "Spider-Man.Homecoming.2017.2160p.BluRay.REMUX.HEVC.mkv", + duplicate_hash, + ) + secondary["updated_at"] = 2 + + async def fetch_rows(media_id): + return [primary] if media_id == manager.media_only_id else [secondary] + + with patch.object(manager, "_fetch_cached_rows", side_effect=fetch_rows): + await manager.get_cached_torrents() + + self.assertIn(duplicate_hash, manager.torrents) + self.assertEqual( + manager.torrents[duplicate_hash]["updatedAt"], + secondary["updated_at"], + ) + self.assertFalse(manager.primary_cached) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_orchestration.py b/tests/test_orchestration.py index 32291605..4b66e2e5 100644 --- a/tests/test_orchestration.py +++ b/tests/test_orchestration.py @@ -196,7 +196,7 @@ async def test_corrupt_cached_parse_does_not_discard_valid_peer(self): media_type="movie", media_full_id="tt123", media_only_id="tt123", - title="Title", + title="Valid", year=2024, year_end=None, season=None,