Aaron Boxer
202 lines
7.1 KiB
Python
202 lines
7.1 KiB
Python
import unittest
|
|
from unittest.mock import MagicMock, patch
|
|
import numpy as np
|
|
|
|
|
|
class TestSpeakerDiarizer(unittest.TestCase):
|
|
"""Tests for SpeakerDiarizer with mocked embedding model."""
|
|
|
|
def _make_diarizer(self, **kwargs):
|
|
from whisper_live.diarization import SpeakerDiarizer
|
|
|
|
d = SpeakerDiarizer(**kwargs)
|
|
# Mock the embedding model to return deterministic embeddings
|
|
d._model = MagicMock()
|
|
return d
|
|
|
|
def _set_embedding(self, diarizer, embedding):
|
|
"""Configure mock model to return a specific embedding."""
|
|
emb = np.array(embedding, dtype=np.float32)
|
|
emb = emb / np.linalg.norm(emb)
|
|
diarizer._model.return_value = emb
|
|
|
|
def test_first_speaker_creates_new(self):
|
|
d = self._make_diarizer()
|
|
self._set_embedding(d, [1.0, 0.0, 0.0])
|
|
audio = np.zeros(16000, dtype=np.float32) # 1 second of audio
|
|
speaker = d.identify_speaker(audio)
|
|
self.assertEqual(speaker, "SPEAKER_00")
|
|
self.assertEqual(len(d.speakers), 1)
|
|
|
|
def test_same_speaker_matches(self):
|
|
d = self._make_diarizer(similarity_threshold=0.8)
|
|
self._set_embedding(d, [1.0, 0.0, 0.0])
|
|
audio = np.zeros(16000, dtype=np.float32)
|
|
d.identify_speaker(audio) # SPEAKER_00
|
|
# Same embedding should match
|
|
self._set_embedding(d, [0.99, 0.01, 0.0])
|
|
speaker = d.identify_speaker(audio)
|
|
self.assertEqual(speaker, "SPEAKER_00")
|
|
self.assertEqual(len(d.speakers), 1)
|
|
|
|
def test_different_speaker_creates_new(self):
|
|
d = self._make_diarizer(similarity_threshold=0.8)
|
|
self._set_embedding(d, [1.0, 0.0, 0.0])
|
|
audio = np.zeros(16000, dtype=np.float32)
|
|
d.identify_speaker(audio) # SPEAKER_00
|
|
|
|
# Very different embedding
|
|
self._set_embedding(d, [0.0, 1.0, 0.0])
|
|
speaker = d.identify_speaker(audio)
|
|
self.assertEqual(speaker, "SPEAKER_01")
|
|
self.assertEqual(len(d.speakers), 2)
|
|
|
|
def test_max_speakers_limit(self):
|
|
d = self._make_diarizer(similarity_threshold=0.95, max_speakers=2)
|
|
audio = np.zeros(16000, dtype=np.float32)
|
|
|
|
self._set_embedding(d, [1.0, 0.0, 0.0])
|
|
d.identify_speaker(audio) # SPEAKER_00
|
|
self._set_embedding(d, [0.0, 1.0, 0.0])
|
|
d.identify_speaker(audio) # SPEAKER_01
|
|
|
|
# Third distinct speaker should be assigned to closest existing
|
|
self._set_embedding(d, [0.0, 0.0, 1.0])
|
|
speaker = d.identify_speaker(audio)
|
|
self.assertIn(speaker, ["SPEAKER_00", "SPEAKER_01"])
|
|
self.assertEqual(len(d.speakers), 2)
|
|
|
|
def test_short_audio_returns_none(self):
|
|
d = self._make_diarizer()
|
|
# Less than 0.3 seconds
|
|
audio = np.zeros(3000, dtype=np.float32)
|
|
speaker = d.identify_speaker(audio)
|
|
self.assertIsNone(speaker)
|
|
|
|
def test_reset_clears_state(self):
|
|
d = self._make_diarizer()
|
|
self._set_embedding(d, [1.0, 0.0, 0.0])
|
|
audio = np.zeros(16000, dtype=np.float32)
|
|
d.identify_speaker(audio)
|
|
self.assertEqual(len(d.speakers), 1)
|
|
d.reset()
|
|
self.assertEqual(len(d.speakers), 0)
|
|
self.assertEqual(d._speaker_count, 0)
|
|
|
|
def test_enroll_speaker_uses_known_name(self):
|
|
d = self._make_diarizer(similarity_threshold=0.8)
|
|
self._set_embedding(d, [1.0, 0.0, 0.0])
|
|
audio = np.zeros(16000, dtype=np.float32)
|
|
self.assertTrue(d.enroll_speaker("Alice", audio))
|
|
|
|
self._set_embedding(d, [0.99, 0.01, 0.0])
|
|
speaker = d.identify_speaker(audio)
|
|
self.assertEqual(speaker, "Alice")
|
|
|
|
def test_speaker_names_label_new_speakers(self):
|
|
d = self._make_diarizer(speaker_names=["Alice"])
|
|
self._set_embedding(d, [1.0, 0.0, 0.0])
|
|
audio = np.zeros(16000, dtype=np.float32)
|
|
speaker = d.identify_speaker(audio)
|
|
self.assertEqual(speaker, "Alice")
|
|
|
|
def test_import_error_without_pyannote(self):
|
|
from whisper_live.diarization import SpeakerDiarizer
|
|
|
|
d = SpeakerDiarizer()
|
|
with patch.dict("sys.modules", {"pyannote": None, "pyannote.audio": None}):
|
|
with self.assertRaises(ImportError):
|
|
d._load_model()
|
|
|
|
|
|
class TestDiarizationInBase(unittest.TestCase):
|
|
"""Test diarization integration in ServeClientBase."""
|
|
|
|
def _make_client(self, diarization=None):
|
|
from whisper_live.backend.base import ServeClientBase
|
|
|
|
class ConcreteClient(ServeClientBase):
|
|
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
|
|
|
|
ws = MagicMock()
|
|
return ConcreteClient(
|
|
client_uid="test-uid", websocket=ws, diarization=diarization
|
|
)
|
|
|
|
def test_no_diarization_by_default(self):
|
|
client = self._make_client()
|
|
self.assertIsNone(client.diarization)
|
|
|
|
def test_format_segment_with_speaker(self):
|
|
client = self._make_client()
|
|
seg = client.format_segment(0.0, 1.0, "hello", speaker="SPEAKER_00")
|
|
self.assertEqual(seg["speaker"], "SPEAKER_00")
|
|
|
|
def test_format_segment_without_speaker(self):
|
|
client = self._make_client()
|
|
seg = client.format_segment(0.0, 1.0, "hello")
|
|
self.assertNotIn("speaker", seg)
|
|
|
|
def test_identify_speaker_disabled(self):
|
|
client = self._make_client(diarization=None)
|
|
seg = MagicMock()
|
|
seg.start = 0.0
|
|
seg.end = 1.0
|
|
result = client._identify_speaker(seg)
|
|
self.assertIsNone(result)
|
|
|
|
def test_identify_speaker_calls_diarizer(self):
|
|
mock_diarizer = MagicMock()
|
|
mock_diarizer.identify_speaker.return_value = "SPEAKER_01"
|
|
client = self._make_client(diarization=mock_diarizer)
|
|
# Set up audio buffer
|
|
client.frames_np = np.zeros(48000, dtype=np.float32)
|
|
client.frames_offset = 0.0
|
|
client.timestamp_offset = 0.0
|
|
seg = MagicMock()
|
|
seg.start = 0.5
|
|
seg.end = 1.5
|
|
result = client._identify_speaker(seg)
|
|
self.assertEqual(result, "SPEAKER_01")
|
|
mock_diarizer.identify_speaker.assert_called_once()
|
|
|
|
|
|
class TestRestDiarizationHelpers(unittest.TestCase):
|
|
def test_normalize_form_list_accepts_repeated_or_comma_separated_values(self):
|
|
from whisper_live.server import TranscriptionServer
|
|
|
|
self.assertEqual(
|
|
TranscriptionServer._normalize_form_list(["Alice,Bob", "Carol"]),
|
|
["Alice", "Bob", "Carol"],
|
|
)
|
|
|
|
def test_speaker_labels_for_segments(self):
|
|
from whisper_live.server import TranscriptionServer
|
|
|
|
segment = MagicMock()
|
|
segment.start = 0.0
|
|
segment.end = 1.0
|
|
diarizer = MagicMock()
|
|
diarizer.identify_speaker.return_value = "Alice"
|
|
|
|
speakers = TranscriptionServer._speaker_labels_for_segments(
|
|
[segment],
|
|
np.zeros(16000, dtype=np.float32),
|
|
diarizer,
|
|
)
|
|
|
|
self.assertEqual(speakers, {0: "Alice"})
|
|
diarizer.identify_speaker.assert_called_once()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|