Rename test script file

tests/{transcription.py → test_transcription.py}

@@ -1,7 +1,9 @@
 from modules.whisper.whisper_factory import WhisperFactory
 from modules.whisper.whisper_parameter import WhisperValues
+from modules.utils.paths import WEBUI_DIR
 from test_config import *
 
+import requests
 import pytest
 import gradio as gr
 import os
@@ -9,39 +11,63 @@ import os
 
 @pytest.mark.parametrize("whisper_type", ["whisper", "faster-whisper", "insanely_fast_whisper"])
 def test_transcribe(whisper_type: str):
-    audio_path = os.path.join("test.wav")
+    audio_path_dir = os.path.join(WEBUI_DIR, "tests")
+    audio_path = os.path.join(audio_path_dir, "jfk.wav")
     if not os.path.exists(audio_path):
-        download_file(TEST_FILE_DOWNLOAD_URL, audio_path)
+        download_file(TEST_FILE_DOWNLOAD_URL, audio_path_dir)
 
     whisper_inferencer = WhisperFactory.create_whisper_inference(
         whisper_type=whisper_type,
     )
-
     print("Device : ", whisper_inferencer.device)
 
     hparams = WhisperValues(
         model_size=TEST_WHISPER_MODEL,
     ).as_list()
 
-    whisper_inferencer.transcribe_file(
-        files=[audio_path],
-        progress=gr.Progress(),
+    subtitle_str, file_path = whisper_inferencer.transcribe_file(
+        [audio_path],
+        None,
+        "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_youtube(
-        youtube_link=TEST_YOUTUBE_URL,
-        progress=gr.Progress(),
+        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(
-        mic_audio=audio_path,
-        progress=gr.Progress(),
+        audio_path,
+        "SRT",
+        False,
+        gr.Progress(),
         *hparams,
     )
+    assert isinstance(subtitle_str, str) and subtitle_str
+    assert isinstance(file_path[0], str) and file_path
+
+
+def download_file(url, save_dir):
+    if not os.path.exists(save_dir):
+        os.makedirs(save_dir)
+
+    file_name = url.split("/")[-1]
+    file_path = os.path.join(save_dir, file_name)
+
+    response = requests.get(url)
 
+    with open(file_path, "wb") as file:
+        file.write(response.content)
 
-def download_file(url: str, path: str):
-    if not os.path.exists(path):
-        os.system(f"wget {url} -O {path}")
+    print(f"File downloaded to: {file_path}")