Add the ability to use YouTube API

This is to circumvent YouTube blocking.

2 files changed, 4 insertions, 4 deletions

diff --git a/tests/obtain_vid_info.py b/tests/obtain_vid_info.py
index 701104d..2bdf330 100644
--- a/tests/obtain_vid_info.py
+++ b/tests/obtain_vid_info.py
@@ -6,12 +6,12 @@ from .utils.vid_url_to_html import get_vid_html_from_url
 class TestObtainVidInfo(TestCase):
     def test_obtain_vid_duration_from_shorts(self) -> None:
         url = "https://www.youtube.com/shorts/iD1Z7ccGyhk"
-        self.assertEqual(60, obtain_vid_duration(url, html=get_vid_html_from_url(url)))
+        self.assertEqual(60, obtain_vid_duration(url, '', html=get_vid_html_from_url(url)))
 
     def test_obtain_vid_duration_from_videos(self) -> None:
         url = "https://www.youtube.com/watch?v=WI4U1SVIO3I"
-        self.assertEqual(8*60+11, obtain_vid_duration(url, html=get_vid_html_from_url(url)))
+        self.assertEqual(8*60+11, obtain_vid_duration(url, '', html=get_vid_html_from_url(url)))
 
     def test_obtain_vid_duration_from_videos_with_params(self) -> None:
         url = "https://www.youtube.com/watch?v=k7RM-ot2NWY&list=PLZHQObOWTQDPD3MizzM2xVFitgF8hE_ab&index=2&pp=iAQB"
-        self.assertEqual(9*60+59, obtain_vid_duration(url, html=get_vid_html_from_url(url)))
+        self.assertEqual(9*60+59, obtain_vid_duration(url, '', html=get_vid_html_from_url(url)))
diff --git a/tests/utils/vid_url_to_html.py b/tests/utils/vid_url_to_html.py
index 60f6ccc..91fcb40 100644
--- a/tests/utils/vid_url_to_html.py
+++ b/tests/utils/vid_url_to_html.py
@@ -18,5 +18,5 @@ def get_vid_html_from_url(url: str) -> str:
     with open(f'tests/data/video@{extract_vid_id(url)}.html', 'r') as file:
         return file.read()
 
-def obtain_vid_duration(url:str, html: str = '') -> int:
+def obtain_vid_duration(url:str, vid_id:str, html:str='', api_key:str='') -> int:
     return get_random_vid_duration(url)