Merge pull request #371 from jhj0517/fix/tests

Update pytube bot detection exception

tests/test_config.py

@@ -1,4 +1,7 @@
+import functools
+
 from modules.utils.paths import *
+from modules.utils.youtube_manager import *
 
 import os
 import torch
@@ -13,5 +16,21 @@
 TEST_SUBTITLE_VTT_PATH = os.path.join(WEBUI_DIR, "tests", "test_vtt.vtt")
 
 
+@functools.lru_cache
 def is_cuda_available():
     return torch.cuda.is_available()
+
+
+@functools.lru_cache
+def is_pytube_detected_bot(url: str = TEST_YOUTUBE_URL):
+    try:
+        yt_temp_path = os.path.join("modules", "yt_tmp.wav")
+        if os.path.exists(yt_temp_path):
+            return False
+        yt = get_ytdata(url)
+        audio = get_ytaudio(yt)
+        return False
+    except Exception as e:
+        print(f"Pytube has detected as a bot: {e}")
+        return True
+

tests/test_transcription.py

@@ -66,15 +66,16 @@ def test_transcribe(
     assert isinstance(subtitle_str, str) and subtitle_str
     assert isinstance(file_path[0], str) and file_path
 
-    whisper_inferencer.transcribe_youtube(
-        TEST_YOUTUBE_URL,
-        "SRT",
-        False,
-        gr.Progress(),
-        *hparams,
-    )
-    assert isinstance(subtitle_str, str) and subtitle_str
-    assert isinstance(file_path[0], str) and file_path
+    if not is_pytube_detected_bot():
+        whisper_inferencer.transcribe_youtube(
+            TEST_YOUTUBE_URL,
+            "SRT",
+            False,
+            gr.Progress(),
+            *hparams,
+        )
+        assert isinstance(subtitle_str, str) and subtitle_str
+        assert isinstance(file_path[0], str) and file_path
 
     whisper_inferencer.transcribe_mic(
         audio_path,