From 0a7933bbba5f0d9744e2a9807c12e8519ecdfa40 Mon Sep 17 00:00:00 2001
From: ducklet <ducklet@noreply.code.dumpr.org>
Date: Sat, 25 May 2024 18:23:13 +0200
Subject: [PATCH] fix: loading multiple oscars per movie

---
 tests/test_web.py                    | 11 +++++++++--
 unwind/cli/import_wikidata_oscars.py | 10 ++++++----
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/tests/test_web.py b/tests/test_web.py
index f9ddbd7..05affd7 100644
--- a/tests/test_web.py
+++ b/tests/test_web.py
@@ -78,7 +78,10 @@ async def test_get_ratings_for_group_with_awards(
     award3 = models.Award(
         movie_id=movie1.id, category="oscars", details='{"name":"Best Visual Effects"}'
     )
-    await db.add(conn, award1, award2, award3)
+    award4 = models.Award(
+        movie_id=movie1.id, category="oscars", details='{"name":"Best Picture"}'
+    )
+    await db.add(conn, award1, award2, award3, award4)
 
     rating = models.Rating(
         movie_id=movie1.id, user_id=user.id, score=66, rating_date=datetime.now(tz=UTC)
@@ -94,7 +97,11 @@ async def test_get_ratings_for_group_with_awards(
         "original_title": movie1.original_title,
         "user_scores": [rating.score],
         "year": movie1.release_year,
-        "awards": ["imdb-top-250:23", "oscars:Best Visual Effects"],
+        "awards": [
+            "imdb-top-250:23",
+            "oscars:Best Picture",
+            "oscars:Best Visual Effects",
+        ],
     }
 
     resp = unauthorized_client.get(path)
diff --git a/unwind/cli/import_wikidata_oscars.py b/unwind/cli/import_wikidata_oscars.py
index 17a9918..cafc341 100644
--- a/unwind/cli/import_wikidata_oscars.py
+++ b/unwind/cli/import_wikidata_oscars.py
@@ -75,7 +75,9 @@ async def main(args: argparse.Namespace) -> None:
 
     json_file: Path = args.json_file
 
-    awards = dict(load_awards(json_file))
+    awards: dict[types.ImdbMovieId, list[models.Award]] = {}
+    for imdb_id, award in load_awards(json_file):
+        awards.setdefault(imdb_id, []).append(award)
     async with db.new_connection() as conn:
         imdb_ids = list(awards)
         available = await db.get_movie_ids(conn, imdb_ids)
@@ -91,9 +93,9 @@ async def main(args: argparse.Namespace) -> None:
         await remove_all_oscars(conn)
 
         for imdb_id, unwind_id in available.items():
-            award = awards[imdb_id]
-            award.movie_id = unwind_id
-            await db.add(conn, award)
+            for award in awards[imdb_id]:
+                award.movie_id = unwind_id
+                await db.add(conn, award)
 
     log.info(f"✨ Imported {len(available)} oscars.")