Skip to content
Open
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions comet/services/filtering.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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
Expand Down
45 changes: 36 additions & 9 deletions comet/services/orchestration.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import asyncio
import time

from pydantic import ValidationError
from RTN import DefaultRanking, ParsedData

from comet.core.execution import get_executor
Expand All @@ -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
Expand Down Expand Up @@ -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 = {}
Expand All @@ -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(
Expand All @@ -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
Expand Down Expand Up @@ -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"]
Expand Down
140 changes: 140 additions & 0 deletions tests/test_cached_title_revalidation.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
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.assertFalse(manager.primary_cached)


if __name__ == "__main__":
unittest.main()
2 changes: 1 addition & 1 deletion tests/test_orchestration.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down