From 85f37aeb2594b05e0edf798132fa5ba198ac98ce Mon Sep 17 00:00:00 2001
From: jhj0517 <97279763+jhj0517@users.noreply.github.com>
Date: Fri, 1 Nov 2024 14:02:05 +0900
Subject: [PATCH] Update pytube bot detection exception

---
 tests/test_config.py        | 19 +++++++++++++++++++
 tests/test_transcription.py | 19 ++++++++++---------
 2 files changed, 29 insertions(+), 9 deletions(-)

diff --git a/tests/test_config.py b/tests/test_config.py
index 0020eca..ba52892 100644
--- a/tests/test_config.py
+++ b/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
+
diff --git a/tests/test_transcription.py b/tests/test_transcription.py
index 3353782..e9591d0 100644
--- a/tests/test_transcription.py
+++ b/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,