58 lines
2.1 KiB
Python
58 lines
2.1 KiB
Python
|
from unittest.mock import MagicMock, patch
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
from whisper_local.transcriber import Transcriber
|
||
|
|
|
||
|
|
|
||
|
|
class TestTranscriber:
|
||
|
|
@patch("whisper_local.transcriber.WhisperModel")
|
||
|
|
def test_init_loads_model(self, mock_model_class):
|
||
|
|
t = Transcriber(model_name="small", compute_type="int8", language="de")
|
||
|
|
mock_model_class.assert_called_once_with("small", compute_type="int8")
|
||
|
|
assert t.language == "de"
|
||
|
|
|
||
|
|
@patch("whisper_local.transcriber.WhisperModel")
|
||
|
|
def test_transcribe_returns_text(self, mock_model_class):
|
||
|
|
mock_model = MagicMock()
|
||
|
|
mock_segment = MagicMock()
|
||
|
|
mock_segment.text = " Hallo Welt "
|
||
|
|
mock_model.transcribe.return_value = ([mock_segment], None)
|
||
|
|
mock_model_class.return_value = mock_model
|
||
|
|
|
||
|
|
t = Transcriber(model_name="small", compute_type="int8", language="de")
|
||
|
|
audio = np.zeros(16000, dtype=np.float32)
|
||
|
|
result = t.transcribe(audio)
|
||
|
|
|
||
|
|
assert result == "Hallo Welt"
|
||
|
|
mock_model.transcribe.assert_called_once_with(audio, language="de")
|
||
|
|
|
||
|
|
@patch("whisper_local.transcriber.WhisperModel")
|
||
|
|
def test_transcribe_empty_segments(self, mock_model_class):
|
||
|
|
mock_model = MagicMock()
|
||
|
|
mock_model.transcribe.return_value = ([], None)
|
||
|
|
mock_model_class.return_value = mock_model
|
||
|
|
|
||
|
|
t = Transcriber(model_name="small", compute_type="int8", language="de")
|
||
|
|
audio = np.zeros(16000, dtype=np.float32)
|
||
|
|
result = t.transcribe(audio)
|
||
|
|
|
||
|
|
assert result == ""
|
||
|
|
|
||
|
|
@patch("whisper_local.transcriber.WhisperModel")
|
||
|
|
def test_transcribe_multiple_segments(self, mock_model_class):
|
||
|
|
mock_model = MagicMock()
|
||
|
|
seg1 = MagicMock()
|
||
|
|
seg1.text = " Hallo "
|
||
|
|
seg2 = MagicMock()
|
||
|
|
seg2.text = " Welt "
|
||
|
|
mock_model.transcribe.return_value = ([seg1, seg2], None)
|
||
|
|
mock_model_class.return_value = mock_model
|
||
|
|
|
||
|
|
t = Transcriber(model_name="small", compute_type="int8", language="de")
|
||
|
|
audio = np.zeros(16000, dtype=np.float32)
|
||
|
|
result = t.transcribe(audio)
|
||
|
|
|
||
|
|
assert result == "Hallo Welt"
|