Merge pull request #436 from boxerab/testing
CI: expand test suite coverage
This commit is contained in:
Vineet Suryan
@@ -0,0 +1,5 @@
|
||||
[pytest]
|
||||
testpaths = tests
|
||||
python_files = test_*.py
|
||||
python_classes = Test*
|
||||
python_functions = test_*
|
||||
@@ -0,0 +1,387 @@
|
||||
import json
|
||||
import queue
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
|
||||
from whisper_live.backend.base import ServeClientBase
|
||||
|
||||
|
||||
class ConcreteServeClient(ServeClientBase):
|
||||
"""Concrete subclass for testing the abstract base class."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self.language = "en"
|
||||
|
||||
def transcribe_audio(self, input_sample):
|
||||
return None
|
||||
|
||||
def handle_transcription_output(self, result, duration):
|
||||
pass
|
||||
|
||||
|
||||
class TestServeClientBaseInit(unittest.TestCase):
|
||||
def test_default_values(self):
|
||||
ws = MagicMock()
|
||||
client = ConcreteServeClient(client_uid="test-uid", websocket=ws)
|
||||
self.assertEqual(client.client_uid, "test-uid")
|
||||
self.assertEqual(client.send_last_n_segments, 10)
|
||||
self.assertAlmostEqual(client.no_speech_thresh, 0.45)
|
||||
self.assertFalse(client.clip_audio)
|
||||
self.assertEqual(client.same_output_threshold, 10)
|
||||
self.assertIsNone(client.frames_np)
|
||||
self.assertAlmostEqual(client.timestamp_offset, 0.0)
|
||||
self.assertFalse(client.exit)
|
||||
self.assertEqual(client.transcript, [])
|
||||
|
||||
def test_custom_values(self):
|
||||
ws = MagicMock()
|
||||
q = queue.Queue()
|
||||
client = ConcreteServeClient(
|
||||
client_uid="uid2",
|
||||
websocket=ws,
|
||||
send_last_n_segments=5,
|
||||
no_speech_thresh=0.6,
|
||||
clip_audio=True,
|
||||
same_output_threshold=20,
|
||||
translation_queue=q,
|
||||
)
|
||||
self.assertEqual(client.send_last_n_segments, 5)
|
||||
self.assertAlmostEqual(client.no_speech_thresh, 0.6)
|
||||
self.assertTrue(client.clip_audio)
|
||||
self.assertEqual(client.same_output_threshold, 20)
|
||||
self.assertIs(client.translation_queue, q)
|
||||
|
||||
|
||||
class TestAddFrames(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
|
||||
|
||||
def test_first_frame_initializes_buffer(self):
|
||||
frame = np.array([0.1, 0.2, 0.3], dtype=np.float32)
|
||||
self.client.add_frames(frame)
|
||||
np.testing.assert_array_equal(self.client.frames_np, frame)
|
||||
|
||||
def test_subsequent_frames_concatenated(self):
|
||||
frame1 = np.array([0.1, 0.2], dtype=np.float32)
|
||||
frame2 = np.array([0.3, 0.4], dtype=np.float32)
|
||||
self.client.add_frames(frame1)
|
||||
self.client.add_frames(frame2)
|
||||
expected = np.array([0.1, 0.2, 0.3, 0.4], dtype=np.float32)
|
||||
np.testing.assert_array_equal(self.client.frames_np, expected)
|
||||
|
||||
def test_buffer_trimmed_at_45_seconds(self):
|
||||
# 45 seconds + 1 sample at 16kHz = 720001 samples
|
||||
self.client.frames_np = np.zeros(45 * 16000 + 1, dtype=np.float32)
|
||||
self.client.add_frames(np.array([1.0], dtype=np.float32))
|
||||
# after trimming 30s, buffer should be ~15s + 1 original + 1 new
|
||||
expected_len = (45 * 16000 + 1) - (30 * 16000) + 1
|
||||
self.assertEqual(self.client.frames_np.shape[0], expected_len)
|
||||
self.assertAlmostEqual(self.client.frames_offset, 30.0)
|
||||
|
||||
def test_timestamp_offset_updated_on_trim(self):
|
||||
self.client.frames_np = np.zeros(45 * 16000 + 1, dtype=np.float32)
|
||||
self.client.timestamp_offset = 5.0 # behind frames_offset after trim
|
||||
self.client.add_frames(np.array([1.0], dtype=np.float32))
|
||||
# timestamp_offset should be bumped to at least frames_offset
|
||||
self.assertGreaterEqual(self.client.timestamp_offset, self.client.frames_offset)
|
||||
|
||||
|
||||
class TestAddFramesThreadSafety(unittest.TestCase):
|
||||
def test_concurrent_add_frames(self):
|
||||
ws = MagicMock()
|
||||
client = ConcreteServeClient(client_uid="test", websocket=ws)
|
||||
errors = []
|
||||
|
||||
def add_many():
|
||||
try:
|
||||
for _ in range(100):
|
||||
client.add_frames(np.random.randn(160).astype(np.float32))
|
||||
except Exception as e:
|
||||
errors.append(e)
|
||||
|
||||
threads = [threading.Thread(target=add_many) for _ in range(4)]
|
||||
for t in threads:
|
||||
t.start()
|
||||
for t in threads:
|
||||
t.join()
|
||||
|
||||
self.assertEqual(errors, [])
|
||||
self.assertIsNotNone(client.frames_np)
|
||||
|
||||
|
||||
class TestGetAudioChunkForProcessing(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
|
||||
|
||||
def test_empty_buffer_returns_empty(self):
|
||||
self.client.frames_np = np.array([], dtype=np.float32)
|
||||
chunk, duration = self.client.get_audio_chunk_for_processing()
|
||||
self.assertEqual(duration, 0.0)
|
||||
self.assertEqual(chunk.shape[0], 0)
|
||||
|
||||
def test_full_buffer_no_offset(self):
|
||||
audio = np.random.randn(16000).astype(np.float32) # 1 second
|
||||
self.client.frames_np = audio
|
||||
chunk, duration = self.client.get_audio_chunk_for_processing()
|
||||
self.assertAlmostEqual(duration, 1.0)
|
||||
np.testing.assert_array_equal(chunk, audio)
|
||||
|
||||
def test_with_offset(self):
|
||||
audio = np.random.randn(32000).astype(np.float32) # 2 seconds
|
||||
self.client.frames_np = audio
|
||||
self.client.timestamp_offset = 1.0 # skip first second
|
||||
chunk, duration = self.client.get_audio_chunk_for_processing()
|
||||
self.assertAlmostEqual(duration, 1.0)
|
||||
self.assertEqual(chunk.shape[0], 16000)
|
||||
|
||||
|
||||
class TestClipAudioIfNoValidSegment(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(
|
||||
client_uid="test", websocket=self.ws, clip_audio=True
|
||||
)
|
||||
|
||||
def test_clips_when_chunk_exceeds_25s(self):
|
||||
# 30 seconds of audio with no valid segments
|
||||
self.client.frames_np = np.zeros(30 * 16000, dtype=np.float32)
|
||||
self.client.timestamp_offset = 0.0
|
||||
self.client.frames_offset = 0.0
|
||||
self.client.clip_audio_if_no_valid_segment()
|
||||
# offset should have advanced to leave ~5s of remaining audio
|
||||
expected_offset = (30 * 16000 / 16000) - 5
|
||||
self.assertAlmostEqual(self.client.timestamp_offset, expected_offset, places=1)
|
||||
|
||||
def test_no_clip_when_short(self):
|
||||
self.client.frames_np = np.zeros(10 * 16000, dtype=np.float32)
|
||||
self.client.timestamp_offset = 0.0
|
||||
self.client.frames_offset = 0.0
|
||||
self.client.clip_audio_if_no_valid_segment()
|
||||
self.assertAlmostEqual(self.client.timestamp_offset, 0.0)
|
||||
|
||||
|
||||
class TestPrepareSegments(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(
|
||||
client_uid="test", websocket=self.ws, send_last_n_segments=3
|
||||
)
|
||||
|
||||
def test_empty_transcript_no_last(self):
|
||||
segments = self.client.prepare_segments()
|
||||
self.assertEqual(segments, [])
|
||||
|
||||
def test_empty_transcript_with_last(self):
|
||||
last = {"start": "0.000", "end": "1.000", "text": "hello", "completed": False}
|
||||
segments = self.client.prepare_segments(last_segment=last)
|
||||
self.assertEqual(len(segments), 1)
|
||||
self.assertEqual(segments[0]["text"], "hello")
|
||||
|
||||
def test_fewer_than_n_segments(self):
|
||||
self.client.transcript = [
|
||||
{"start": "0.000", "end": "1.000", "text": "a", "completed": True},
|
||||
{"start": "1.000", "end": "2.000", "text": "b", "completed": True},
|
||||
]
|
||||
segments = self.client.prepare_segments()
|
||||
self.assertEqual(len(segments), 2)
|
||||
|
||||
def test_more_than_n_segments_truncated(self):
|
||||
self.client.transcript = [
|
||||
{"start": f"{i}.000", "end": f"{i+1}.000", "text": f"seg{i}", "completed": True}
|
||||
for i in range(10)
|
||||
]
|
||||
segments = self.client.prepare_segments()
|
||||
self.assertEqual(len(segments), 3)
|
||||
self.assertEqual(segments[0]["text"], "seg7")
|
||||
|
||||
def test_last_segment_appended(self):
|
||||
self.client.transcript = [
|
||||
{"start": "0.000", "end": "1.000", "text": "a", "completed": True},
|
||||
]
|
||||
last = {"start": "1.000", "end": "2.000", "text": "in progress", "completed": False}
|
||||
segments = self.client.prepare_segments(last_segment=last)
|
||||
self.assertEqual(len(segments), 2)
|
||||
self.assertEqual(segments[-1]["text"], "in progress")
|
||||
|
||||
|
||||
class TestFormatSegment(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
|
||||
|
||||
def test_format(self):
|
||||
seg = self.client.format_segment(1.234, 5.678, "hello world", completed=True)
|
||||
self.assertEqual(seg["start"], "1.234")
|
||||
self.assertEqual(seg["end"], "5.678")
|
||||
self.assertEqual(seg["text"], "hello world")
|
||||
self.assertTrue(seg["completed"])
|
||||
|
||||
def test_format_not_completed(self):
|
||||
seg = self.client.format_segment(0.0, 1.0, "text")
|
||||
self.assertFalse(seg["completed"])
|
||||
|
||||
|
||||
class TestSendTranscriptionToClient(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(client_uid="test-uid", websocket=self.ws)
|
||||
|
||||
def test_sends_json(self):
|
||||
segments = [{"start": "0.000", "end": "1.000", "text": "hi", "completed": True}]
|
||||
self.client.send_transcription_to_client(segments)
|
||||
self.ws.send.assert_called_once()
|
||||
sent = json.loads(self.ws.send.call_args[0][0])
|
||||
self.assertEqual(sent["uid"], "test-uid")
|
||||
self.assertEqual(len(sent["segments"]), 1)
|
||||
|
||||
def test_send_failure_logged_not_raised(self):
|
||||
self.ws.send.side_effect = ConnectionError("broken pipe")
|
||||
# should not raise
|
||||
self.client.send_transcription_to_client([])
|
||||
|
||||
|
||||
class TestDisconnect(unittest.TestCase):
|
||||
def test_sends_disconnect_message(self):
|
||||
ws = MagicMock()
|
||||
client = ConcreteServeClient(client_uid="uid1", websocket=ws)
|
||||
client.disconnect()
|
||||
sent = json.loads(ws.send.call_args[0][0])
|
||||
self.assertEqual(sent["uid"], "uid1")
|
||||
self.assertEqual(sent["message"], "DISCONNECT")
|
||||
|
||||
|
||||
class TestCleanup(unittest.TestCase):
|
||||
def test_sets_exit_flag(self):
|
||||
ws = MagicMock()
|
||||
client = ConcreteServeClient(client_uid="uid1", websocket=ws)
|
||||
self.assertFalse(client.exit)
|
||||
client.cleanup()
|
||||
self.assertTrue(client.exit)
|
||||
|
||||
|
||||
class TestUpdateSegments(unittest.TestCase):
|
||||
"""Tests for the core update_segments() logic."""
|
||||
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(
|
||||
client_uid="test",
|
||||
websocket=self.ws,
|
||||
no_speech_thresh=0.45,
|
||||
same_output_threshold=3,
|
||||
)
|
||||
self.client.frames_np = np.zeros(16000 * 5, dtype=np.float32)
|
||||
|
||||
def _make_segment(self, start, end, text, no_speech_prob=0.0):
|
||||
seg = MagicMock()
|
||||
seg.start = start
|
||||
seg.end = end
|
||||
seg.text = text
|
||||
seg.no_speech_prob = no_speech_prob
|
||||
return seg
|
||||
|
||||
def test_single_segment_becomes_last(self):
|
||||
segs = [self._make_segment(0.0, 1.0, " hello")]
|
||||
last = self.client.update_segments(segs, duration=2.0)
|
||||
self.assertIsNotNone(last)
|
||||
self.assertIn("hello", last["text"])
|
||||
self.assertFalse(last["completed"])
|
||||
self.assertEqual(len(self.client.transcript), 0)
|
||||
|
||||
def test_multiple_segments_completes_all_but_last(self):
|
||||
segs = [
|
||||
self._make_segment(0.0, 1.0, " first"),
|
||||
self._make_segment(1.0, 2.0, " second"),
|
||||
]
|
||||
last = self.client.update_segments(segs, duration=3.0)
|
||||
self.assertEqual(len(self.client.transcript), 1)
|
||||
self.assertTrue(self.client.transcript[0]["completed"])
|
||||
self.assertIn("first", self.client.transcript[0]["text"])
|
||||
self.assertIsNotNone(last)
|
||||
self.assertIn("second", last["text"])
|
||||
|
||||
def test_high_no_speech_prob_skipped(self):
|
||||
segs = [
|
||||
self._make_segment(0.0, 1.0, " noise", no_speech_prob=0.9),
|
||||
self._make_segment(1.0, 2.0, " also noise", no_speech_prob=0.9),
|
||||
]
|
||||
last = self.client.update_segments(segs, duration=3.0)
|
||||
self.assertEqual(len(self.client.transcript), 0)
|
||||
self.assertIsNone(last)
|
||||
|
||||
def test_segment_with_start_gte_end_skipped(self):
|
||||
segs = [
|
||||
self._make_segment(1.0, 0.5, " backwards"),
|
||||
self._make_segment(1.5, 2.0, " normal"),
|
||||
]
|
||||
last = self.client.update_segments(segs, duration=3.0)
|
||||
self.assertEqual(len(self.client.transcript), 0)
|
||||
self.assertIsNotNone(last)
|
||||
|
||||
def test_repeated_output_triggers_completion(self):
|
||||
seg = self._make_segment(0.0, 1.0, " repeated")
|
||||
for _ in range(self.client.same_output_threshold + 2):
|
||||
last = self.client.update_segments([seg], duration=2.0)
|
||||
# after enough repeats, should be added to transcript
|
||||
self.assertTrue(len(self.client.transcript) >= 1)
|
||||
|
||||
def test_translation_queue_receives_completed(self):
|
||||
q = queue.Queue()
|
||||
self.client.translation_queue = q
|
||||
segs = [
|
||||
self._make_segment(0.0, 1.0, " first"),
|
||||
self._make_segment(1.0, 2.0, " second"),
|
||||
]
|
||||
self.client.update_segments(segs, duration=3.0)
|
||||
self.assertFalse(q.empty())
|
||||
item = q.get_nowait()
|
||||
self.assertIn("first", item["text"])
|
||||
|
||||
def test_timestamp_offset_advances(self):
|
||||
segs = [
|
||||
self._make_segment(0.0, 1.0, " first"),
|
||||
self._make_segment(1.0, 2.0, " second"),
|
||||
]
|
||||
self.client.update_segments(segs, duration=3.0)
|
||||
self.assertGreater(self.client.timestamp_offset, 0.0)
|
||||
|
||||
|
||||
class TestGetSegmentHelpers(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.ws = MagicMock()
|
||||
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
|
||||
|
||||
def test_get_segment_no_speech_prob_attr(self):
|
||||
seg = MagicMock()
|
||||
seg.no_speech_prob = 0.3
|
||||
self.assertAlmostEqual(self.client.get_segment_no_speech_prob(seg), 0.3)
|
||||
|
||||
def test_get_segment_no_speech_prob_fallback(self):
|
||||
seg = MagicMock(spec=[]) # no attributes
|
||||
self.assertEqual(self.client.get_segment_no_speech_prob(seg), 0)
|
||||
|
||||
def test_get_segment_start_uses_start(self):
|
||||
seg = MagicMock()
|
||||
seg.start = 1.5
|
||||
self.assertAlmostEqual(self.client.get_segment_start(seg), 1.5)
|
||||
|
||||
def test_get_segment_end_uses_end(self):
|
||||
seg = MagicMock()
|
||||
seg.end = 3.0
|
||||
self.assertAlmostEqual(self.client.get_segment_end(seg), 3.0)
|
||||
|
||||
def test_get_segment_start_fallback_to_start_ts(self):
|
||||
seg = MagicMock(spec=["start_ts"])
|
||||
seg.start_ts = 2.0
|
||||
self.assertAlmostEqual(self.client.get_segment_start(seg), 2.0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,257 @@
|
||||
import json
|
||||
import time
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock, PropertyMock
|
||||
|
||||
from whisper_live.client import Client, TranscriptionTeeClient
|
||||
|
||||
|
||||
class TestClientStatusMessages(unittest.TestCase):
|
||||
"""Tests for Client.handle_status_messages() and on_message() branches."""
|
||||
|
||||
@patch("whisper_live.client.websocket.WebSocketApp")
|
||||
@patch("whisper_live.client.pyaudio.PyAudio")
|
||||
def setUp(self, mock_pyaudio, mock_websocket):
|
||||
mock_pyaudio.return_value.open.return_value = MagicMock()
|
||||
self.client = Client(host="localhost", port=9090, lang="en")
|
||||
|
||||
def tearDown(self):
|
||||
self.client.close_websocket()
|
||||
|
||||
def test_wait_status(self):
|
||||
msg = {"uid": self.client.uid, "status": "WAIT", "message": 5.0}
|
||||
self.client.handle_status_messages(msg)
|
||||
self.assertTrue(self.client.waiting)
|
||||
|
||||
def test_error_status(self):
|
||||
msg = {"uid": self.client.uid, "status": "ERROR", "message": "model not found"}
|
||||
self.client.handle_status_messages(msg)
|
||||
self.assertTrue(self.client.server_error)
|
||||
|
||||
def test_warning_status_no_side_effects(self):
|
||||
msg = {"uid": self.client.uid, "status": "WARNING", "message": "fallback backend"}
|
||||
self.client.handle_status_messages(msg)
|
||||
self.assertFalse(self.client.server_error)
|
||||
self.assertFalse(self.client.waiting)
|
||||
|
||||
def test_on_message_wrong_uid_ignored(self):
|
||||
msg = json.dumps({"uid": "wrong-uid", "segments": [{"start": 0, "end": 1, "text": "hi", "completed": True}]})
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
self.assertEqual(len(self.client.transcript), 0)
|
||||
|
||||
def test_on_message_disconnect(self):
|
||||
self.client.recording = True
|
||||
msg = json.dumps({"uid": self.client.uid, "message": "DISCONNECT"})
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
self.assertFalse(self.client.recording)
|
||||
|
||||
def test_on_message_server_ready(self):
|
||||
msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"message": "SERVER_READY",
|
||||
"backend": "faster_whisper",
|
||||
})
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
self.assertTrue(self.client.recording)
|
||||
self.assertEqual(self.client.server_backend, "faster_whisper")
|
||||
|
||||
def test_on_message_language_detection(self):
|
||||
msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"language": "fr",
|
||||
"language_prob": 0.95,
|
||||
})
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
self.assertEqual(self.client.language, "fr")
|
||||
|
||||
|
||||
class TestClientTranslationFlow(unittest.TestCase):
|
||||
"""Tests for the translation-related client functionality."""
|
||||
|
||||
@patch("whisper_live.client.websocket.WebSocketApp")
|
||||
@patch("whisper_live.client.pyaudio.PyAudio")
|
||||
def setUp(self, mock_pyaudio, mock_websocket):
|
||||
mock_pyaudio.return_value.open.return_value = MagicMock()
|
||||
self.client = Client(
|
||||
host="localhost",
|
||||
port=9090,
|
||||
lang="en",
|
||||
enable_translation=True,
|
||||
target_language="es",
|
||||
)
|
||||
# simulate SERVER_READY so server_backend is set
|
||||
ready_msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"message": "SERVER_READY",
|
||||
"backend": "faster_whisper",
|
||||
})
|
||||
self.client.on_message(MagicMock(), ready_msg)
|
||||
|
||||
def tearDown(self):
|
||||
self.client.close_websocket()
|
||||
|
||||
def test_on_open_includes_translation_fields(self):
|
||||
mock_ws = MagicMock()
|
||||
self.client.on_open(mock_ws)
|
||||
sent = json.loads(mock_ws.send.call_args[0][0])
|
||||
self.assertTrue(sent["enable_translation"])
|
||||
self.assertEqual(sent["target_language"], "es")
|
||||
|
||||
def test_translated_segments_processed(self):
|
||||
msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"translated_segments": [
|
||||
{"start": "0.000", "end": "1.000", "text": "Hola mundo", "completed": True},
|
||||
],
|
||||
})
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
self.assertEqual(len(self.client.translated_transcript), 1)
|
||||
self.assertEqual(self.client.translated_transcript[0]["text"], "Hola mundo")
|
||||
|
||||
def test_translation_callback_invoked(self):
|
||||
callback = MagicMock()
|
||||
self.client.translation_callback = callback
|
||||
msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"translated_segments": [
|
||||
{"start": "0.000", "end": "1.000", "text": "Hola", "completed": True},
|
||||
],
|
||||
})
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
callback.assert_called_once()
|
||||
|
||||
def test_translation_callback_exception_handled(self):
|
||||
callback = MagicMock(side_effect=RuntimeError("callback broke"))
|
||||
self.client.translation_callback = callback
|
||||
msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"translated_segments": [
|
||||
{"start": "0.000", "end": "1.000", "text": "Hola", "completed": True},
|
||||
],
|
||||
})
|
||||
# should not raise
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
|
||||
|
||||
class TestClientTranscriptionCallback(unittest.TestCase):
|
||||
"""Tests for the transcription callback feature."""
|
||||
|
||||
@patch("whisper_live.client.websocket.WebSocketApp")
|
||||
@patch("whisper_live.client.pyaudio.PyAudio")
|
||||
def setUp(self, mock_pyaudio, mock_websocket):
|
||||
mock_pyaudio.return_value.open.return_value = MagicMock()
|
||||
self.callback = MagicMock()
|
||||
self.client = Client(
|
||||
host="localhost",
|
||||
port=9090,
|
||||
lang="en",
|
||||
transcription_callback=self.callback,
|
||||
)
|
||||
ready_msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"message": "SERVER_READY",
|
||||
"backend": "faster_whisper",
|
||||
})
|
||||
self.client.on_message(MagicMock(), ready_msg)
|
||||
|
||||
def tearDown(self):
|
||||
self.client.close_websocket()
|
||||
|
||||
def test_callback_receives_text_and_segments(self):
|
||||
msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"segments": [
|
||||
{"start": "0.000", "end": "1.000", "text": "Hello", "completed": True},
|
||||
],
|
||||
})
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
self.callback.assert_called_once()
|
||||
text_arg, segments_arg = self.callback.call_args[0]
|
||||
self.assertIn("Hello", text_arg)
|
||||
self.assertIsInstance(segments_arg, list)
|
||||
|
||||
def test_callback_exception_does_not_crash(self):
|
||||
self.callback.side_effect = ValueError("boom")
|
||||
msg = json.dumps({
|
||||
"uid": self.client.uid,
|
||||
"segments": [
|
||||
{"start": "0.000", "end": "1.000", "text": "Test", "completed": True},
|
||||
],
|
||||
})
|
||||
# should not raise
|
||||
self.client.on_message(MagicMock(), msg)
|
||||
|
||||
|
||||
class TestClientSrtWriting(unittest.TestCase):
|
||||
"""Tests for Client.write_srt_file() edge cases."""
|
||||
|
||||
@patch("whisper_live.client.websocket.WebSocketApp")
|
||||
@patch("whisper_live.client.pyaudio.PyAudio")
|
||||
def setUp(self, mock_pyaudio, mock_websocket):
|
||||
mock_pyaudio.return_value.open.return_value = MagicMock()
|
||||
self.client = Client(host="localhost", port=9090, lang="en")
|
||||
self.client.server_backend = "faster_whisper"
|
||||
|
||||
def tearDown(self):
|
||||
self.client.close_websocket()
|
||||
import os
|
||||
for f in ["test_out.srt"]:
|
||||
if os.path.exists(f):
|
||||
os.remove(f)
|
||||
|
||||
def test_write_srt_empty_transcript_with_last_segment(self):
|
||||
self.client.transcript = []
|
||||
self.client.last_segment = {"start": "0.000", "end": "1.000", "text": "final"}
|
||||
self.client.write_srt_file("test_out.srt")
|
||||
self.assertEqual(len(self.client.transcript), 1)
|
||||
self.assertEqual(self.client.transcript[0]["text"], "final")
|
||||
|
||||
def test_write_srt_appends_last_segment_if_different(self):
|
||||
self.client.transcript = [{"start": "0.000", "end": "1.000", "text": "first"}]
|
||||
self.client.last_segment = {"start": "1.000", "end": "2.000", "text": "second"}
|
||||
self.client.write_srt_file("test_out.srt")
|
||||
self.assertEqual(len(self.client.transcript), 2)
|
||||
|
||||
def test_write_srt_no_duplicate_last_segment(self):
|
||||
self.client.transcript = [{"start": "0.000", "end": "1.000", "text": "same"}]
|
||||
self.client.last_segment = {"start": "0.000", "end": "1.000", "text": "same"}
|
||||
self.client.write_srt_file("test_out.srt")
|
||||
self.assertEqual(len(self.client.transcript), 1)
|
||||
|
||||
|
||||
class TestWaitBeforeDisconnect(unittest.TestCase):
|
||||
"""Tests for Client.wait_before_disconnect()."""
|
||||
|
||||
@patch("whisper_live.client.websocket.WebSocketApp")
|
||||
@patch("whisper_live.client.pyaudio.PyAudio")
|
||||
def setUp(self, mock_pyaudio, mock_websocket):
|
||||
mock_pyaudio.return_value.open.return_value = MagicMock()
|
||||
self.client = Client(host="localhost", port=9090, lang="en")
|
||||
|
||||
def tearDown(self):
|
||||
self.client.close_websocket()
|
||||
|
||||
def test_raises_if_no_response(self):
|
||||
self.client.last_response_received = None
|
||||
with self.assertRaises(AssertionError):
|
||||
self.client.wait_before_disconnect()
|
||||
|
||||
def test_returns_immediately_if_timeout_elapsed(self):
|
||||
self.client.last_response_received = time.time() - 100
|
||||
self.client.disconnect_if_no_response_for = 15
|
||||
start = time.time()
|
||||
self.client.wait_before_disconnect()
|
||||
elapsed = time.time() - start
|
||||
self.assertLess(elapsed, 1.0)
|
||||
|
||||
|
||||
class TestTeeClientEdgeCases(unittest.TestCase):
|
||||
"""Edge cases for TranscriptionTeeClient."""
|
||||
|
||||
def test_empty_clients_raises(self):
|
||||
with self.assertRaises(Exception):
|
||||
TranscriptionTeeClient([])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,217 @@
|
||||
import json
|
||||
import time
|
||||
import unittest
|
||||
from unittest import mock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from whisper_live.server import TranscriptionServer, BackendType, ClientManager
|
||||
|
||||
|
||||
class TestClientManagerAddRemove(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.cm = ClientManager(max_clients=2, max_connection_time=60)
|
||||
|
||||
def test_add_and_get_client(self):
|
||||
ws = MagicMock()
|
||||
client = MagicMock()
|
||||
self.cm.add_client(ws, client)
|
||||
self.assertIs(self.cm.get_client(ws), client)
|
||||
|
||||
def test_get_nonexistent_client(self):
|
||||
ws = MagicMock()
|
||||
self.assertFalse(self.cm.get_client(ws))
|
||||
|
||||
def test_remove_client_calls_cleanup(self):
|
||||
ws = MagicMock()
|
||||
client = MagicMock()
|
||||
self.cm.add_client(ws, client)
|
||||
self.cm.remove_client(ws)
|
||||
client.cleanup.assert_called_once()
|
||||
self.assertNotIn(ws, self.cm.clients)
|
||||
self.assertNotIn(ws, self.cm.start_times)
|
||||
|
||||
def test_remove_nonexistent_client_no_error(self):
|
||||
ws = MagicMock()
|
||||
self.cm.remove_client(ws) # should not raise
|
||||
|
||||
|
||||
class TestClientManagerServerFull(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.cm = ClientManager(max_clients=1, max_connection_time=60)
|
||||
|
||||
def test_not_full_returns_false(self):
|
||||
ws = MagicMock()
|
||||
options = {"uid": "test"}
|
||||
self.assertFalse(self.cm.is_server_full(ws, options))
|
||||
|
||||
def test_full_sends_wait_and_returns_true(self):
|
||||
ws1 = MagicMock()
|
||||
self.cm.add_client(ws1, MagicMock())
|
||||
|
||||
ws2 = MagicMock()
|
||||
options = {"uid": "new-client"}
|
||||
self.assertTrue(self.cm.is_server_full(ws2, options))
|
||||
ws2.send.assert_called_once()
|
||||
sent = json.loads(ws2.send.call_args[0][0])
|
||||
self.assertEqual(sent["status"], "WAIT")
|
||||
self.assertEqual(sent["uid"], "new-client")
|
||||
|
||||
|
||||
class TestClientManagerTimeout(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.cm = ClientManager(max_clients=4, max_connection_time=10)
|
||||
|
||||
def test_not_timed_out(self):
|
||||
ws = MagicMock()
|
||||
client = MagicMock()
|
||||
self.cm.add_client(ws, client)
|
||||
self.assertFalse(self.cm.is_client_timeout(ws))
|
||||
|
||||
def test_timed_out(self):
|
||||
ws = MagicMock()
|
||||
client = MagicMock()
|
||||
self.cm.add_client(ws, client)
|
||||
self.cm.start_times[ws] = time.time() - 20
|
||||
self.assertTrue(self.cm.is_client_timeout(ws))
|
||||
client.disconnect.assert_called_once()
|
||||
|
||||
|
||||
class TestClientManagerGetWaitTime(unittest.TestCase):
|
||||
def test_no_clients_returns_zero(self):
|
||||
cm = ClientManager(max_clients=4, max_connection_time=600)
|
||||
self.assertEqual(cm.get_wait_time(), 0)
|
||||
|
||||
def test_single_client_wait_time(self):
|
||||
cm = ClientManager(max_clients=4, max_connection_time=600)
|
||||
ws = MagicMock()
|
||||
cm.add_client(ws, MagicMock())
|
||||
cm.start_times[ws] = time.time() - 300
|
||||
wait = cm.get_wait_time()
|
||||
self.assertAlmostEqual(wait, 5.0, places=0)
|
||||
|
||||
def test_multiple_clients_returns_minimum(self):
|
||||
cm = ClientManager(max_clients=4, max_connection_time=600)
|
||||
ws1, ws2 = MagicMock(), MagicMock()
|
||||
cm.add_client(ws1, MagicMock())
|
||||
cm.add_client(ws2, MagicMock())
|
||||
cm.start_times[ws1] = time.time() - 100
|
||||
cm.start_times[ws2] = time.time() - 500
|
||||
wait = cm.get_wait_time()
|
||||
# ws2 has 100s remaining = ~1.67 minutes
|
||||
self.assertAlmostEqual(wait, 100 / 60, places=0)
|
||||
|
||||
|
||||
class TestBackendType(unittest.TestCase):
|
||||
def test_valid_types(self):
|
||||
valid = BackendType.valid_types()
|
||||
self.assertIn("faster_whisper", valid)
|
||||
self.assertIn("tensorrt", valid)
|
||||
self.assertIn("openvino", valid)
|
||||
|
||||
def test_is_valid(self):
|
||||
self.assertTrue(BackendType.is_valid("faster_whisper"))
|
||||
self.assertFalse(BackendType.is_valid("nonexistent"))
|
||||
|
||||
def test_type_checks(self):
|
||||
self.assertTrue(BackendType.FASTER_WHISPER.is_faster_whisper())
|
||||
self.assertFalse(BackendType.FASTER_WHISPER.is_tensorrt())
|
||||
self.assertTrue(BackendType.TENSORRT.is_tensorrt())
|
||||
self.assertTrue(BackendType.OPENVINO.is_openvino())
|
||||
|
||||
def test_enum_from_string(self):
|
||||
bt = BackendType("faster_whisper")
|
||||
self.assertEqual(bt, BackendType.FASTER_WHISPER)
|
||||
|
||||
def test_invalid_enum_raises(self):
|
||||
with self.assertRaises(ValueError):
|
||||
BackendType("invalid_backend")
|
||||
|
||||
|
||||
class TestTranscriptionServerInit(unittest.TestCase):
|
||||
def test_defaults(self):
|
||||
server = TranscriptionServer()
|
||||
self.assertIsNone(server.client_manager)
|
||||
self.assertTrue(server.use_vad)
|
||||
self.assertFalse(server.single_model)
|
||||
self.assertIsNone(server.batch_config)
|
||||
|
||||
def test_run_invalid_backend_raises(self):
|
||||
server = TranscriptionServer()
|
||||
with self.assertRaises(ValueError):
|
||||
server.run(host="localhost", port=9090, backend="nonexistent")
|
||||
|
||||
def test_run_invalid_trt_path_raises(self):
|
||||
server = TranscriptionServer()
|
||||
with self.assertRaises(ValueError):
|
||||
server.run(
|
||||
host="localhost",
|
||||
port=9090,
|
||||
backend="tensorrt",
|
||||
whisper_tensorrt_path="/nonexistent/path",
|
||||
)
|
||||
|
||||
|
||||
class TestTranscriptionServerGetAudio(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.server = TranscriptionServer()
|
||||
|
||||
def test_end_of_audio_returns_false(self):
|
||||
ws = MagicMock()
|
||||
ws.recv.return_value = b"END_OF_AUDIO"
|
||||
result = self.server.get_audio_from_websocket(ws)
|
||||
self.assertFalse(result)
|
||||
|
||||
def test_valid_audio_returns_numpy(self):
|
||||
import numpy as np
|
||||
ws = MagicMock()
|
||||
audio = np.array([0.1, 0.2, 0.3], dtype=np.float32)
|
||||
ws.recv.return_value = audio.tobytes()
|
||||
result = self.server.get_audio_from_websocket(ws)
|
||||
np.testing.assert_array_almost_equal(result, audio)
|
||||
|
||||
|
||||
class TestTranscriptionServerHandleNewConnection(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.server = TranscriptionServer()
|
||||
self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
|
||||
self.server.cache_path = "~/.cache/whisper-live/"
|
||||
self.server.backend = BackendType.FASTER_WHISPER
|
||||
|
||||
@mock.patch("websockets.WebSocketCommonProtocol")
|
||||
def test_invalid_json_returns_false(self, mock_ws):
|
||||
mock_ws.recv.return_value = "not valid json {{"
|
||||
result = self.server.handle_new_connection(mock_ws, None, None, False)
|
||||
self.assertFalse(result)
|
||||
|
||||
@mock.patch("websockets.WebSocketCommonProtocol")
|
||||
def test_server_full_returns_false(self, mock_ws):
|
||||
# Fill server
|
||||
for i in range(4):
|
||||
self.server.client_manager.add_client(MagicMock(), MagicMock())
|
||||
|
||||
mock_ws.recv.return_value = json.dumps({
|
||||
"uid": "test",
|
||||
"language": "en",
|
||||
"task": "transcribe",
|
||||
"model": "tiny.en",
|
||||
})
|
||||
result = self.server.handle_new_connection(mock_ws, None, None, False)
|
||||
self.assertFalse(result)
|
||||
|
||||
|
||||
class TestTranscriptionServerCleanup(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.server = TranscriptionServer()
|
||||
self.server.client_manager = ClientManager(max_clients=4, max_connection_time=600)
|
||||
|
||||
def test_cleanup_removes_client(self):
|
||||
ws = MagicMock()
|
||||
client = MagicMock()
|
||||
self.server.client_manager.add_client(ws, client)
|
||||
self.server.cleanup(ws)
|
||||
self.assertNotIn(ws, self.server.client_manager.clients)
|
||||
client.cleanup.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,134 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from io import StringIO
|
||||
from unittest.mock import patch
|
||||
|
||||
from whisper_live.utils import format_time, create_srt_file, print_transcript
|
||||
|
||||
|
||||
class TestFormatTime(unittest.TestCase):
|
||||
def test_zero(self):
|
||||
self.assertEqual(format_time(0), "00:00:00,000")
|
||||
|
||||
def test_seconds_only(self):
|
||||
self.assertEqual(format_time(5.0), "00:00:05,000")
|
||||
|
||||
def test_fractional_seconds(self):
|
||||
self.assertEqual(format_time(1.5), "00:00:01,500")
|
||||
|
||||
def test_minutes(self):
|
||||
self.assertEqual(format_time(65.0), "00:01:05,000")
|
||||
|
||||
def test_hours(self):
|
||||
self.assertEqual(format_time(3661.123), "01:01:01,123")
|
||||
|
||||
def test_millisecond_precision(self):
|
||||
self.assertEqual(format_time(0.001), "00:00:00,001")
|
||||
|
||||
def test_large_value(self):
|
||||
# float precision: int((86399.999 - 86399) * 1000) may be 998 or 999
|
||||
result = format_time(86399.999)
|
||||
self.assertIn(result, ("23:59:59,998", "23:59:59,999"))
|
||||
|
||||
def test_rounding_edge(self):
|
||||
result = format_time(0.9999)
|
||||
# 0.9999 -> int(s%60)=0, milliseconds=int(0.9999*1000)=999
|
||||
self.assertEqual(result, "00:00:00,999")
|
||||
|
||||
|
||||
class TestCreateSrtFile(unittest.TestCase):
|
||||
def test_single_segment(self):
|
||||
segments = [{"start": "0.000", "end": "1.500", "text": "Hello world"}]
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
|
||||
path = f.name
|
||||
try:
|
||||
create_srt_file(segments, path)
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
self.assertIn("1\n", content)
|
||||
self.assertIn("00:00:00,000 --> 00:00:01,500", content)
|
||||
self.assertIn("Hello world", content)
|
||||
finally:
|
||||
os.remove(path)
|
||||
|
||||
def test_multiple_segments(self):
|
||||
segments = [
|
||||
{"start": "0.000", "end": "1.000", "text": "First"},
|
||||
{"start": "1.000", "end": "2.500", "text": "Second"},
|
||||
{"start": "2.500", "end": "4.000", "text": "Third"},
|
||||
]
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
|
||||
path = f.name
|
||||
try:
|
||||
create_srt_file(segments, path)
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
self.assertIn("1\n", content)
|
||||
self.assertIn("2\n", content)
|
||||
self.assertIn("3\n", content)
|
||||
self.assertIn("First", content)
|
||||
self.assertIn("Third", content)
|
||||
finally:
|
||||
os.remove(path)
|
||||
|
||||
def test_empty_segments(self):
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
|
||||
path = f.name
|
||||
try:
|
||||
create_srt_file([], path)
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
self.assertEqual(content, "")
|
||||
finally:
|
||||
os.remove(path)
|
||||
|
||||
def test_unicode_text(self):
|
||||
segments = [{"start": "0.000", "end": "1.000", "text": "日本語テスト"}]
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".srt", delete=False) as f:
|
||||
path = f.name
|
||||
try:
|
||||
create_srt_file(segments, path)
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
self.assertIn("日本語テスト", content)
|
||||
finally:
|
||||
os.remove(path)
|
||||
|
||||
|
||||
class TestPrintTranscript(unittest.TestCase):
|
||||
@patch("sys.stdout", new_callable=StringIO)
|
||||
def test_print_plain_text(self, mock_stdout):
|
||||
text = ["Hello", " world"]
|
||||
print_transcript(text)
|
||||
output = mock_stdout.getvalue()
|
||||
self.assertIn("Hello world", output)
|
||||
|
||||
@patch("sys.stdout", new_callable=StringIO)
|
||||
def test_print_with_timestamps(self, mock_stdout):
|
||||
text = [
|
||||
{"start": 0.0, "end": 1.0, "text": "Hello"},
|
||||
{"start": 1.0, "end": 2.0, "text": "world"},
|
||||
]
|
||||
print_transcript(text, timestamps=True)
|
||||
output = mock_stdout.getvalue()
|
||||
self.assertIn("[0.0 -> 1.0]", output)
|
||||
self.assertIn("Hello", output)
|
||||
|
||||
@patch("sys.stdout", new_callable=StringIO)
|
||||
def test_print_translated(self, mock_stdout):
|
||||
text = ["Bonjour", "le monde"]
|
||||
print_transcript(text, translated=True)
|
||||
output = mock_stdout.getvalue()
|
||||
self.assertIn("Bonjour le monde", output)
|
||||
|
||||
@patch("sys.stdout", new_callable=StringIO)
|
||||
def test_print_empty(self, mock_stdout):
|
||||
print_transcript([])
|
||||
output = mock_stdout.getvalue()
|
||||
# empty text joined is empty string, should not crash
|
||||
self.assertEqual(output.strip(), "")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,131 @@
|
||||
import unittest
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from whisper_live.vad import VoiceActivityDetection, VoiceActivityDetector
|
||||
|
||||
|
||||
class TestVoiceActivityDetectionValidation(unittest.TestCase):
|
||||
"""Tests for VoiceActivityDetection input validation without requiring the ONNX model."""
|
||||
|
||||
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
|
||||
def setUp(self):
|
||||
self.vad = VoiceActivityDetection()
|
||||
self.vad.sample_rates = [8000, 16000]
|
||||
|
||||
def test_1d_input_unsqueezed(self):
|
||||
x = torch.randn(512)
|
||||
x_out, sr_out = self.vad._validate_input(x, 16000)
|
||||
self.assertEqual(x_out.dim(), 2)
|
||||
self.assertEqual(sr_out, 16000)
|
||||
|
||||
def test_3d_input_raises(self):
|
||||
x = torch.randn(1, 1, 512)
|
||||
with self.assertRaises(ValueError):
|
||||
self.vad._validate_input(x, 16000)
|
||||
|
||||
def test_unsupported_sample_rate_raises(self):
|
||||
x = torch.randn(1, 512)
|
||||
with self.assertRaises(ValueError):
|
||||
self.vad._validate_input(x, 44100)
|
||||
|
||||
def test_too_short_audio_raises(self):
|
||||
x = torch.randn(1, 1)
|
||||
with self.assertRaises(ValueError):
|
||||
self.vad._validate_input(x, 16000)
|
||||
|
||||
def test_downsample_multiple_of_16k(self):
|
||||
x = torch.randn(1, 512 * 3)
|
||||
x_out, sr_out = self.vad._validate_input(x, 48000)
|
||||
self.assertEqual(sr_out, 16000)
|
||||
self.assertEqual(x_out.shape[1], 512)
|
||||
|
||||
|
||||
class TestVoiceActivityDetectionStateReset(unittest.TestCase):
|
||||
"""Tests for VoiceActivityDetection.reset_states()."""
|
||||
|
||||
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
|
||||
def setUp(self):
|
||||
self.vad = VoiceActivityDetection()
|
||||
|
||||
def test_reset_creates_correct_shapes(self):
|
||||
self.vad.reset_states(batch_size=4)
|
||||
self.assertEqual(self.vad._state.shape, (2, 4, 128))
|
||||
self.assertEqual(self.vad._context.shape[0], 0)
|
||||
self.assertEqual(self.vad._last_sr, 0)
|
||||
self.assertEqual(self.vad._last_batch_size, 0)
|
||||
|
||||
def test_reset_default_batch_size(self):
|
||||
self.vad.reset_states()
|
||||
self.assertEqual(self.vad._state.shape, (2, 1, 128))
|
||||
|
||||
|
||||
class TestVoiceActivityDetectionDownload(unittest.TestCase):
|
||||
"""Tests for the model download function."""
|
||||
|
||||
@patch("os.path.exists", return_value=True)
|
||||
def test_skips_download_if_exists(self, mock_exists):
|
||||
path = VoiceActivityDetection.download()
|
||||
self.assertTrue(path.endswith("silero_vad.onnx"))
|
||||
|
||||
@patch("os.path.exists", return_value=False)
|
||||
@patch("subprocess.run")
|
||||
@patch("os.makedirs")
|
||||
def test_downloads_if_missing(self, mock_makedirs, mock_run, mock_exists):
|
||||
path = VoiceActivityDetection.download()
|
||||
mock_run.assert_called_once()
|
||||
self.assertIn("silero_vad.onnx", path)
|
||||
|
||||
@patch("os.path.exists", return_value=False)
|
||||
@patch("subprocess.run", side_effect=Exception("wget not found"))
|
||||
@patch("os.makedirs")
|
||||
def test_handles_download_failure(self, mock_makedirs, mock_run, mock_exists):
|
||||
# should not raise, just prints an error
|
||||
with self.assertRaises(Exception):
|
||||
VoiceActivityDetection.download()
|
||||
|
||||
|
||||
class TestVoiceActivityDetectorThreshold(unittest.TestCase):
|
||||
"""Tests for VoiceActivityDetector threshold behavior."""
|
||||
|
||||
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
|
||||
def test_above_threshold_returns_true(self):
|
||||
detector = VoiceActivityDetector.__new__(VoiceActivityDetector)
|
||||
detector.model = VoiceActivityDetection()
|
||||
detector.threshold = 0.5
|
||||
detector.frame_rate = 16000
|
||||
|
||||
mock_probs = torch.tensor([[0.9, 0.8, 0.7]])
|
||||
with patch.object(detector.model, "audio_forward", return_value=mock_probs):
|
||||
result = detector(np.random.randn(16000).astype(np.float32))
|
||||
self.assertTrue(result)
|
||||
|
||||
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
|
||||
def test_below_threshold_returns_false(self):
|
||||
detector = VoiceActivityDetector.__new__(VoiceActivityDetector)
|
||||
detector.model = VoiceActivityDetection()
|
||||
detector.threshold = 0.5
|
||||
detector.frame_rate = 16000
|
||||
|
||||
mock_probs = torch.tensor([[0.1, 0.2, 0.3]])
|
||||
with patch.object(detector.model, "audio_forward", return_value=mock_probs):
|
||||
result = detector(np.random.randn(16000).astype(np.float32))
|
||||
self.assertFalse(result)
|
||||
|
||||
@patch.object(VoiceActivityDetection, "__init__", lambda self, **kw: None)
|
||||
def test_custom_threshold(self):
|
||||
detector = VoiceActivityDetector.__new__(VoiceActivityDetector)
|
||||
detector.model = VoiceActivityDetection()
|
||||
detector.threshold = 0.95
|
||||
detector.frame_rate = 16000
|
||||
|
||||
mock_probs = torch.tensor([[0.9]])
|
||||
with patch.object(detector.model, "audio_forward", return_value=mock_probs):
|
||||
result = detector(np.random.randn(16000).astype(np.float32))
|
||||
self.assertFalse(result)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user