mirror of
https://github.com/chidiwilliams/buzz.git
synced 2026-03-14 14:45:46 +01:00
111 lines
4 KiB
Python
111 lines
4 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,
|
|
DEFAULT_WHISPER_TEMPERATURE,
|
|
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,
|
|
)
|
|
|
|
|
|
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()
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks={},
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory="/path/to/output/folder",
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
temperature=DEFAULT_WHISPER_TEMPERATURE,
|
|
initial_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
default_export_file_name="",
|
|
)
|
|
|
|
shutil.copy("testdata/whisper-french.mp3", 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 == "/path/to/output/folder"
|
|
|
|
def test_should_not_add_task_in_tasks(self, qtbot):
|
|
input_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="/path/to/output/folder",
|
|
model_path="",
|
|
),
|
|
}
|
|
|
|
watcher = TranscriptionTaskFolderWatcher(
|
|
tasks=tasks,
|
|
preferences=FolderWatchPreferences(
|
|
enabled=True,
|
|
input_directory=input_directory,
|
|
output_directory="/path/to/output/folder",
|
|
file_transcription_options=FileTranscriptionPreferences(
|
|
language=None,
|
|
task=Task.TRANSCRIBE,
|
|
model=self.default_model(),
|
|
word_level_timings=False,
|
|
temperature=DEFAULT_WHISPER_TEMPERATURE,
|
|
initial_prompt="",
|
|
output_formats=set(),
|
|
),
|
|
),
|
|
default_export_file_name="",
|
|
)
|
|
|
|
# Ignored because already in tasks
|
|
shutil.copy(
|
|
"testdata/whisper-french.mp3",
|
|
os.path.join(input_directory, "whisper-french.mp3"),
|
|
)
|
|
shutil.copy(
|
|
"testdata/whisper-french.mp3",
|
|
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 == "/path/to/output/folder"
|