From fe9cc7f29fe81806e342141390352fe24cbd4da6 Mon Sep 17 00:00:00 2001
From: Soulsuck24 <79275800+Soulsuck24@users.noreply.github.com>
Date: Sat, 24 Dec 2022 20:52:45 -0500
Subject: [PATCH] Redgifs updates

Update Redgifs regex for further edge case.

Add test for checking ID.
---
 bdfr/site_downloaders/redgifs.py       |  9 +++++++--
 tests/site_downloaders/test_redgifs.py | 13 +++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/bdfr/site_downloaders/redgifs.py b/bdfr/site_downloaders/redgifs.py
index 40d1466..95d23d2 100644
--- a/bdfr/site_downloaders/redgifs.py
+++ b/bdfr/site_downloaders/redgifs.py
@@ -23,13 +23,18 @@ class Redgifs(BaseDownloader):
         return [Resource(self.post, m, Resource.retry_download(m), None) for m in media_urls]
 
     @staticmethod
-    def _get_link(url: str) -> set[str]:
+    def _get_id(url: str) -> str:
         try:
             if url.endswith("/"):
                 url = url.removesuffix("/")
-            redgif_id = re.match(r".*/(.*?)(\..{0,})?$", url).group(1)
+            redgif_id = re.match(r".*/(.*?)(?:\?.*|\..{0,})?$", url).group(1)
         except AttributeError:
             raise SiteDownloaderError(f"Could not extract Redgifs ID from {url}")
+        return redgif_id
+
+    @staticmethod
+    def _get_link(url: str) -> set[str]:
+        redgif_id = Redgifs._get_id(url)
 
         auth_token = json.loads(Redgifs.retrieve_url("https://api.redgifs.com/v2/auth/temporary").text)["token"]
         if not auth_token:
diff --git a/tests/site_downloaders/test_redgifs.py b/tests/site_downloaders/test_redgifs.py
index 37e48e8..fd0e0ed 100644
--- a/tests/site_downloaders/test_redgifs.py
+++ b/tests/site_downloaders/test_redgifs.py
@@ -10,6 +10,19 @@ from bdfr.resource import Resource
 from bdfr.site_downloaders.redgifs import Redgifs
 
 
+@pytest.mark.parametrize(
+    ("test_url", "expected"),
+    (
+        ("https://redgifs.com/watch/frighteningvictorioussalamander", "frighteningvictorioussalamander"),
+        ("https://www.redgifs.com/watch/genuineprivateguillemot/", "genuineprivateguillemot"),
+        ("https://www.redgifs.com/watch/marriedcrushingcob?rel=u%3Akokiri.girl%3Bo%3Arecent", "marriedcrushingcob"),
+    ),
+)
+def test_get_id(test_url: str, expected: str):
+    result = Redgifs._get_id(test_url)
+    assert result == expected
+
+
 @pytest.mark.online
 @pytest.mark.parametrize(
     ("test_url", "expected"),