Aaron Boxer
b1cd51ac8a
these new test cover issues such as thread safety, VAD thresholding, message routing, error handling etc. that weren't covered by existing tests. Mocking is used to avoid dependencies on GPU, ONNX etc.
132 lines
4.9 KiB
Python
132 lines
4.9 KiB
Python
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()
|