From 0f94c98733f191255f5b85f5cb17162b184495cc Mon Sep 17 00:00:00 2001
From: OMEGARAZER <869111+OMEGARAZER@users.noreply.github.com>
Date: Mon, 9 Jan 2023 12:48:24 -0500
Subject: [PATCH] Account for new gallery url

Coverage for gallery urls
---
 bdfr/site_downloaders/imgur.py       |  9 ++--
 tests/site_downloaders/test_imgur.py | 63 +++++++---------------------
 2 files changed, 21 insertions(+), 51 deletions(-)

diff --git a/bdfr/site_downloaders/imgur.py b/bdfr/site_downloaders/imgur.py
index 537ab64..4af62b1 100644
--- a/bdfr/site_downloaders/imgur.py
+++ b/bdfr/site_downloaders/imgur.py
@@ -42,9 +42,12 @@ class Imgur(BaseDownloader):
     @staticmethod
     def _get_data(link: str) -> dict:
         try:
-            imgur_id = re.match(r".*/(.*?)(_d)?(\..{0,})?$", link).group(1)
-            gallery = "a/" if re.search(r".*/(.*?)(gallery/|a/)", link) else ""
-            link = f"https://imgur.com/{gallery}{imgur_id}"
+            if re.search(r".*/(.*?)(gallery/|a/)", link):
+                imgur_id = re.match(r".*/(?:gallery/|a/)(.*?)(?:/.*)?$", link).group(1)
+                link = f"https://imgur.com/a/{imgur_id}"
+            else:
+                imgur_id = re.match(r".*/(.*?)(_d)?(\..{0,})?$", link).group(1)
+                link = f"https://imgur.com/{imgur_id}"
         except AttributeError:
             raise SiteDownloaderError(f"Could not extract Imgur ID from {link}")
 
diff --git a/tests/site_downloaders/test_imgur.py b/tests/site_downloaders/test_imgur.py
index 7f587c9..9f8e6a7 100644
--- a/tests/site_downloaders/test_imgur.py
+++ b/tests/site_downloaders/test_imgur.py
@@ -109,10 +109,7 @@ def test_imgur_extension_validation_bad(test_extension: str):
     ("test_url", "expected_hashes"),
     (
         ("https://imgur.com/a/xWZsDDP", ("f551d6e6b0fef2ce909767338612e31b",)),
-        (
-            "https://imgur.com/gallery/IjJJdlC",
-            ("740b006cf9ec9d6f734b6e8f5130bdab",),
-        ),
+        ("https://imgur.com/gallery/IjJJdlC", ("740b006cf9ec9d6f734b6e8f5130bdab",)),
         (
             "https://imgur.com/a/dcc84Gt",
             (
@@ -130,50 +127,20 @@ def test_imgur_extension_validation_bad(test_extension: str):
                 "fb6c913d721c0bbb96aa65d7f560d385",
             ),
         ),
-        (
-            "https://i.imgur.com/lFJai6i.gifv",
-            ("01a6e79a30bec0e644e5da12365d5071",),
-        ),
-        (
-            "https://i.imgur.com/ywSyILa.gifv?",
-            ("56d4afc32d2966017c38d98568709b45",),
-        ),
-        (
-            "https://imgur.com/ubYwpbk.GIFV",
-            ("d4a774aac1667783f9ed3a1bd02fac0c",),
-        ),
-        (
-            "https://i.imgur.com/j1CNCZY.gifv",
-            ("58e7e6d972058c18b7ecde910ca147e3",),
-        ),
-        (
-            "https://i.imgur.com/uTvtQsw.gifv",
-            ("46c86533aa60fc0e09f2a758513e3ac2",),
-        ),
-        (
-            "https://i.imgur.com/OGeVuAe.giff",
-            ("77389679084d381336f168538793f218",),
-        ),
-        (
-            "https://i.imgur.com/OGeVuAe.gift",
-            ("77389679084d381336f168538793f218",),
-        ),
-        (
-            "https://i.imgur.com/3SKrQfK.jpg?1",
-            ("aa299e181b268578979cad176d1bd1d0",),
-        ),
-        (
-            "https://i.imgur.com/cbivYRW.jpg?3",
-            ("7ec6ceef5380cb163a1d498c359c51fd",),
-        ),
-        (
-            "http://i.imgur.com/s9uXxlq.jpg?5.jpg",
-            ("338de3c23ee21af056b3a7c154e2478f",),
-        ),
-        (
-            "https://i.imgur.com/2TtN68l_d.webp",
-            ("6569ab9ad9fa68d93f6b408f112dd741",),
-        ),
+        ("https://i.imgur.com/lFJai6i.gifv", ("01a6e79a30bec0e644e5da12365d5071",)),
+        ("https://i.imgur.com/ywSyILa.gifv?", ("56d4afc32d2966017c38d98568709b45",)),
+        ("https://imgur.com/ubYwpbk.GIFV", ("d4a774aac1667783f9ed3a1bd02fac0c",)),
+        ("https://i.imgur.com/j1CNCZY.gifv", ("58e7e6d972058c18b7ecde910ca147e3",)),
+        ("https://i.imgur.com/uTvtQsw.gifv", ("46c86533aa60fc0e09f2a758513e3ac2",)),
+        ("https://i.imgur.com/OGeVuAe.giff", ("77389679084d381336f168538793f218",)),
+        ("https://i.imgur.com/OGeVuAe.gift", ("77389679084d381336f168538793f218",)),
+        ("https://i.imgur.com/3SKrQfK.jpg?1", ("aa299e181b268578979cad176d1bd1d0",)),
+        ("https://i.imgur.com/cbivYRW.jpg?3", ("7ec6ceef5380cb163a1d498c359c51fd",)),
+        ("http://i.imgur.com/s9uXxlq.jpg?5.jpg", ("338de3c23ee21af056b3a7c154e2478f",)),
+        ("https://i.imgur.com/2TtN68l_d.webp", ("6569ab9ad9fa68d93f6b408f112dd741",)),
+        ("https://imgur.com/a/1qzfWtY/gifv", ("65fbc7ba5c3ed0e3af47c4feef4d3735",)),
+        ("https://imgur.com/a/1qzfWtY/mp4", ("65fbc7ba5c3ed0e3af47c4feef4d3735",)),
+        ("https://imgur.com/a/1qzfWtY/spqr", ("65fbc7ba5c3ed0e3af47c4feef4d3735",)),
     ),
 )
 def test_find_resources(test_url: str, expected_hashes: list[str]):