Add unit tests for obtain_vid_info

1 files changed, 17 insertions, 0 deletions

diff --git a/tests/obtain_vid_info.py b/tests/obtain_vid_info.py
new file mode 100644
index 0000000..701104d
--- /dev/null
+++ b/tests/obtain_vid_info.py
@@ -0,0 +1,17 @@
+from unittest import TestCase
+
+from components.extractor.obtain_vid_info import obtain_vid_duration
+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)))
+
+    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)))
+
+    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)))