Add unit tests for obtain_vid_info

1 files changed, 18 insertions, 0 deletions

diff --git a/tests/utils/vid_url_to_html.py b/tests/utils/vid_url_to_html.py
new file mode 100644
index 0000000..355e830
--- /dev/null
+++ b/tests/utils/vid_url_to_html.py
@@ -0,0 +1,18 @@
+from urllib.parse import urlparse, parse_qs
+
+from components.extractor.check_url import is_video
+
+def extract_vid_id(url: str) -> str:
+    if not is_video(url):
+        raise Exception(url + " is not a YouTube video URL")
+    parsed_url = urlparse(url)
+    if parsed_url.netloc == 'youtu.be':
+        return parsed_url.path[1:]
+    if parsed_url.path == '/watch':
+        query_params = parse_qs(parsed_url.query)
+        return query_params['v'][0]
+    return parsed_url.path.split('/')[-1]
+
+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()