fix ruff and add yamtrack partial season and episode test

This commit is contained in:
FuzzyGrim
2026-02-01 19:15:24 +01:00
parent 7d4deac23f
commit f27d26d4ec
2 changed files with 72 additions and 24 deletions

View File

@@ -245,27 +245,24 @@ def get_media_metadata(
def search(media_type, query, page, source=None):
"""Search for media based on the query and return the results."""
if media_type == MediaTypes.MANGA.value:
if source == Sources.MANGAUPDATES.value:
response = mangaupdates.search(query, page)
else:
response = mal.search(media_type, query, page)
elif media_type == MediaTypes.ANIME.value:
response = mal.search(media_type, query, page)
elif media_type in (MediaTypes.TV.value, MediaTypes.MOVIE.value):
response = tmdb.search(media_type, query, page)
elif media_type in (MediaTypes.SEASON.value, MediaTypes.EPISODE.value):
response = tmdb.search(MediaTypes.TV.value, query, page)
elif media_type == MediaTypes.GAME.value:
response = igdb.search(query, page)
elif media_type == MediaTypes.BOOK.value:
if source == Sources.OPENLIBRARY.value:
response = openlibrary.search(query, page)
else:
response = hardcover.search(query, page)
elif media_type == MediaTypes.COMIC.value:
response = comicvine.search(query, page)
elif media_type == MediaTypes.BOARDGAME.value:
response = bgg.search(query, page)
return response
search_handlers = {
MediaTypes.MANGA.value: lambda: (
mangaupdates.search(query, page)
if source == Sources.MANGAUPDATES.value
else mal.search(media_type, query, page)
),
MediaTypes.ANIME.value: lambda: mal.search(media_type, query, page),
MediaTypes.TV.value: lambda: tmdb.search(media_type, query, page),
MediaTypes.MOVIE.value: lambda: tmdb.search(media_type, query, page),
MediaTypes.SEASON.value: lambda: tmdb.search(MediaTypes.TV.value, query, page),
MediaTypes.EPISODE.value: lambda: tmdb.search(MediaTypes.TV.value, query, page),
MediaTypes.GAME.value: lambda: igdb.search(query, page),
MediaTypes.BOOK.value: lambda: (
openlibrary.search(query, page)
if source == Sources.OPENLIBRARY.value
else hardcover.search(query, page)
),
MediaTypes.COMIC.value: lambda: comicvine.search(query, page),
MediaTypes.BOARDGAME.value: lambda: bgg.search(query, page),
}
return search_handlers[media_type]()

View File

@@ -135,6 +135,57 @@ class ImportYamtrackPartials(TestCase):
self.assertEqual(Book.objects.filter(user=self.user).count(), 3)
self.assertEqual(Movie.objects.filter(user=self.user).count(), 1)
def test_season_episode_search_by_title(self):
"""Test that seasons and episodes can be resolved by title (no media_id).
This test verifies the fix that allows searching for SEASON and EPISODE
media types by searching for the parent TV show on TMDB. Before the fix,
services.search() didn't handle SEASON/EPISODE types and would fail with:
UnboundLocalError: cannot access local variable 'response'
"""
test_rows = [
# Season with title only (no media_id)
{
"media_id": "",
"source": "",
"media_type": "season",
"title": "Friends",
"image": "",
"season_number": "1",
"episode_number": "",
},
# Episode with title only (no media_id)
{
"media_id": "",
"source": "",
"media_type": "episode",
"title": "Friends",
"image": "",
"season_number": "1",
"episode_number": "1",
},
]
importer = yamtrack.YamtrackImporter(None, self.user, "new")
for row in test_rows:
original_row = row.copy()
importer._handle_missing_metadata(
row,
row["media_type"],
int(row["season_number"]) if row["season_number"] else None,
int(row["episode_number"]) if row["episode_number"] else None,
)
# Verify media_id was resolved from TMDB search
self.assertNotEqual(row["media_id"], original_row["media_id"])
self.assertEqual(str(row["media_id"]), "1668") # Friends TV show ID
self.assertEqual(row["source"], "tmdb")
# Title and image should be populated from TMDB
self.assertNotEqual(row["title"], "")
self.assertNotEqual(row["image"], "")
def test_end_dates(self):
"""Test end dates during import."""
book = Book.objects.filter(user=self.user).first()