mirror of
https://github.com/chidiwilliams/buzz.git
synced 2026-03-14 14:45:46 +01:00
418 lines
15 KiB
Python
418 lines
15 KiB
Python
import os
|
|
import shutil
|
|
from tempfile import mkdtemp
|
|
|
|
from pytestqt.qtbot import QtBot
|
|
|
|
from buzz.model_loader import TranscriptionModel, ModelType
|
|
from buzz.transcriber.transcriber import (
|
|
Task,
|
|
FileTranscriptionTask,
|
|
TranscriptionOptions,
|
|
FileTranscriptionOptions,
|
|
)
|
|
from buzz.widgets.preferences_dialog.models.file_transcription_preferences import (
|
|
FileTranscriptionPreferences,
|
|
)
|
|
from buzz.widgets.preferences_dialog.models.folder_watch_preferences import (
|
|
FolderWatchPreferences,
|
|
)
|
|
from buzz.widgets.transcription_task_folder_watcher import (
|
|
TranscriptionTaskFolderWatcher,
|
|
)
|
|
from tests.audio import test_audio_path
|
|
|
|
|
|
class TestTranscriptionTaskFolderWatcher:
|
|
def default_model(self):
|
|
model_type = next(
|
|
model_type for model_type in ModelType if model_type.is_available()
|
|
)
|
|
return TranscriptionModel(model_type=model_type)
|
|
|
|
def test_should_add_task_not_in_tasks(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
shutil.copy(test_audio_path, input_directory)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.file_path == os.path.join(input_directory, "whisper-french.mp3")
|
|
assert task.source == FileTranscriptionTask.Source.FOLDER_WATCH
|
|
assert task.output_directory == output_directory
|
|
|
|
def test_should_not_add_task_in_tasks(self, qtbot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
tasks = {
|
|
1: FileTranscriptionTask(
|
|
file_path=os.path.join(input_directory, "whisper-french.mp3"),
|
|
transcription_options=TranscriptionOptions(),
|
|
file_transcription_options=FileTranscriptionOptions(file_paths=[]),
|
|
output_directory=output_directory,
|
|
model_path="",
|
|
),
|
|
}
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks=tasks,
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
# Ignored because already in tasks
|
|
shutil.copy(
|
|
test_audio_path,
|
|
os.path.join(input_directory, "whisper-french.mp3"),
|
|
)
|
|
shutil.copy(
|
|
test_audio_path,
|
|
os.path.join(input_directory, "whisper-french2.mp3"),
|
|
)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.file_path == os.path.join(input_directory, "whisper-french2.mp3")
|
|
assert task.source == FileTranscriptionTask.Source.FOLDER_WATCH
|
|
assert task.output_directory == output_directory
|
|
|
|
def test_should_find_tasks_in_subfolders(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
subfolder = os.path.join(input_directory, "subfolder")
|
|
os.makedirs(subfolder)
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
shutil.copy(test_audio_path, subfolder)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.file_path == os.path.join(subfolder, "whisper-french.mp3")
|
|
assert task.output_directory == os.path.join(output_directory, "subfolder")
|
|
assert os.path.isdir(task.output_directory)
|
|
|
|
def test_should_ignore_hidden_directories(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
hidden_folder = os.path.join(input_directory, ".hidden")
|
|
os.makedirs(hidden_folder)
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
# Copy to hidden folder (should be ignored)
|
|
shutil.copy(test_audio_path, hidden_folder)
|
|
# Copy to root folder (should be found)
|
|
shutil.copy(
|
|
test_audio_path,
|
|
os.path.join(input_directory, "visible.mp3"),
|
|
)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.file_path == os.path.join(input_directory, "visible.mp3")
|
|
|
|
def test_should_ignore_non_media_files(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
# Create non-media files (should be ignored)
|
|
with open(os.path.join(input_directory, "test.txt"), "w") as f:
|
|
f.write("test")
|
|
with open(os.path.join(input_directory, "test.json"), "w") as f:
|
|
f.write("{}")
|
|
|
|
# Copy media file (should be found)
|
|
shutil.copy(test_audio_path, input_directory)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.file_path == os.path.join(input_directory, "whisper-french.mp3")
|
|
|
|
def test_should_ignore_temp_conversion_files(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
# Create temp conversion file (should be ignored)
|
|
shutil.copy(
|
|
test_audio_path,
|
|
os.path.join(input_directory, "audio.ogg.wav"),
|
|
)
|
|
|
|
# Copy normal media file (should be found)
|
|
shutil.copy(test_audio_path, input_directory)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.file_path == os.path.join(input_directory, "whisper-french.mp3")
|
|
|
|
def test_should_ignore_extracted_speech_files(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
# Create extracted speech file (should be ignored)
|
|
shutil.copy(
|
|
test_audio_path,
|
|
os.path.join(input_directory, "video_speech.mp3"),
|
|
)
|
|
|
|
# Copy normal media file (should be found)
|
|
shutil.copy(test_audio_path, input_directory)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.file_path == os.path.join(input_directory, "whisper-french.mp3")
|
|
|
|
def test_should_set_delete_source_file_when_preference_enabled(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
delete_processed_files=True,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
shutil.copy(test_audio_path, input_directory)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.delete_source_file is True
|
|
|
|
def test_should_not_set_delete_source_file_when_preference_disabled(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
delete_processed_files=False,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
shutil.copy(test_audio_path, input_directory)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
assert task.delete_source_file is False
|
|
|
|
def test_should_set_original_file_path(self, qtbot: QtBot):
|
|
input_directory = mkdtemp()
|
|
output_directory = mkdtemp()
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory=output_directory,
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
extract_speech=False,
|
|
initial_prompt="",
|
|
enable_llm_translation=False,
|
|
llm_model="",
|
|
llm_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
)
|
|
|
|
shutil.copy(test_audio_path, input_directory)
|
|
|
|
with qtbot.wait_signal(watcher.task_found, timeout=10_000) as blocker:
|
|
pass
|
|
|
|
task: FileTranscriptionTask = blocker.args[0]
|
|
expected_path = os.path.join(input_directory, "whisper-french.mp3")
|
|
assert task.file_path == expected_path
|
|
assert task.original_file_path == expected_path
|