David Maier
610 lines
23 KiB
Python
610 lines
23 KiB
Python
import json
|
|
import queue
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from unittest.mock import MagicMock
|
|
|
|
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 WaitTrackingEvent:
|
|
"""Threading event that records when wait() is entered."""
|
|
|
|
def __init__(self):
|
|
self._event = threading.Event()
|
|
self.wait_started = threading.Event()
|
|
|
|
def wait(self, timeout=None):
|
|
self.wait_started.set()
|
|
return self._event.wait(timeout)
|
|
|
|
def set(self):
|
|
self._event.set()
|
|
|
|
|
|
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 TestSpeechToTextWaitingBehavior(unittest.TestCase):
|
|
"""Tests the first-frame wait behavior in speech_to_text()."""
|
|
|
|
def setUp(self):
|
|
self.ws = MagicMock()
|
|
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
|
|
self.client.frames_ready = WaitTrackingEvent()
|
|
self.transcribe_called = threading.Event()
|
|
self.thread_started = threading.Event()
|
|
self.cpu_used = None
|
|
self.thread = None
|
|
|
|
def tearDown(self):
|
|
if self.thread is not None and self.thread.is_alive():
|
|
self.client.exit = True
|
|
# release wait() directly so a broken cleanup() cannot hang the test process
|
|
self.client.frames_ready.set()
|
|
self.thread.join(timeout=1.0)
|
|
|
|
def _start_speech_thread(self, target=None):
|
|
self.thread = threading.Thread(target=target or self.client.speech_to_text)
|
|
self.thread.start()
|
|
return self.thread
|
|
|
|
def _join_speech_thread(self):
|
|
self.thread.join(timeout=1.0)
|
|
return not self.thread.is_alive()
|
|
|
|
def _transcribe_once(self, input_sample):
|
|
# mark the first processing step after wait and stop the loop
|
|
self.transcribe_called.set()
|
|
self.client.exit = True
|
|
return []
|
|
|
|
def _measure_waiting_cpu(self):
|
|
# measure CPU consumed by speech_to_text loop while it waits for the first frame
|
|
self.thread_started.set()
|
|
start_cpu = time.thread_time()
|
|
self.client.speech_to_text()
|
|
self.cpu_used = time.thread_time() - start_cpu
|
|
|
|
def test_waits_for_first_frame_before_transcribing(self):
|
|
self.client.transcribe_audio = MagicMock(side_effect=self._transcribe_once)
|
|
self._start_speech_thread()
|
|
self.assertTrue(self.client.frames_ready.wait_started.wait(timeout=1.0))
|
|
self.assertFalse(self.transcribe_called.is_set())
|
|
|
|
self.client.add_frames(np.zeros(self.client.RATE, dtype=np.float32))
|
|
|
|
self.assertTrue(self.transcribe_called.wait(timeout=1.0))
|
|
self.assertTrue(self._join_speech_thread())
|
|
|
|
def test_cleanup_unblocks_waiting_thread_without_audio(self):
|
|
self.client.transcribe_audio = MagicMock()
|
|
self._start_speech_thread()
|
|
self.assertTrue(self.client.frames_ready.wait_started.wait(timeout=1.0))
|
|
|
|
self.client.cleanup()
|
|
|
|
self.assertTrue(self._join_speech_thread())
|
|
self.assertTrue(self.client.exit)
|
|
self.client.transcribe_audio.assert_not_called()
|
|
|
|
def test_waiting_for_first_frame_uses_negligible_thread_cpu(self):
|
|
self._start_speech_thread(target=self._measure_waiting_cpu)
|
|
self.assertTrue(self.thread_started.wait(timeout=1.0))
|
|
|
|
time.sleep(0.25)
|
|
self.client.cleanup()
|
|
|
|
self.assertTrue(self._join_speech_thread())
|
|
self.assertIsNotNone(self.cpu_used)
|
|
self.assertLess(self.cpu_used, 0.05)
|
|
|
|
|
|
class TestTrimTranscript(unittest.TestCase):
|
|
def setUp(self):
|
|
self.ws = MagicMock()
|
|
self.client = ConcreteServeClient(client_uid="test", websocket=self.ws)
|
|
|
|
def test_transcript_trimmed_when_over_max(self):
|
|
self.client.transcript = [
|
|
{"start": f"{i}.000", "end": f"{i+1}.000", "text": f"seg{i}", "completed": True}
|
|
for i in range(self.client.MAX_TRANSCRIPT_LENGTH + 100)
|
|
]
|
|
self.client._trim_transcript()
|
|
self.assertEqual(len(self.client.transcript), self.client.MAX_TRANSCRIPT_LENGTH)
|
|
self.assertEqual(self.client.transcript[0]["text"], "seg100")
|
|
|
|
def test_transcript_not_trimmed_when_under_max(self):
|
|
self.client.transcript = [
|
|
{"start": "0.000", "end": "1.000", "text": "a", "completed": True}
|
|
]
|
|
self.client._trim_transcript()
|
|
self.assertEqual(len(self.client.transcript), 1)
|
|
|
|
def test_text_list_trimmed(self):
|
|
self.client.text = ["word"] * (self.client.MAX_TRANSCRIPT_LENGTH + 50)
|
|
self.client._trim_transcript()
|
|
self.assertEqual(len(self.client.text), self.client.MAX_TRANSCRIPT_LENGTH)
|
|
|
|
|
|
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)
|
|
|
|
|
|
class TestWordTimestamps(unittest.TestCase):
|
|
"""Tests for word-level timestamp extraction."""
|
|
|
|
def _make_client(self, word_timestamps=False):
|
|
ws = MagicMock()
|
|
return ConcreteServeClient(
|
|
client_uid="wt-uid", websocket=ws, word_timestamps=word_timestamps
|
|
)
|
|
|
|
def _make_word(self, word, start, end, prob):
|
|
w = MagicMock()
|
|
w.word = word
|
|
w.start = start
|
|
w.end = end
|
|
w.probability = prob
|
|
return w
|
|
|
|
def _make_segment(self, text, start, end, no_speech_prob=0.0, words=None):
|
|
seg = MagicMock()
|
|
seg.text = text
|
|
seg.start = start
|
|
seg.end = end
|
|
seg.no_speech_prob = no_speech_prob
|
|
seg.words = words
|
|
return seg
|
|
|
|
def test_word_timestamps_disabled_by_default(self):
|
|
client = self._make_client()
|
|
self.assertFalse(client.word_timestamps)
|
|
|
|
def test_word_timestamps_enabled(self):
|
|
client = self._make_client(word_timestamps=True)
|
|
self.assertTrue(client.word_timestamps)
|
|
|
|
def test_extract_words_when_disabled(self):
|
|
client = self._make_client(word_timestamps=False)
|
|
seg = self._make_segment("hello", 0.0, 1.0, words=[self._make_word("hello", 0.0, 0.5, 0.99)])
|
|
result = client._extract_words(seg, 0.0)
|
|
self.assertIsNone(result)
|
|
|
|
def test_extract_words_when_enabled(self):
|
|
client = self._make_client(word_timestamps=True)
|
|
words = [
|
|
self._make_word("hello", 0.0, 0.3, 0.95),
|
|
self._make_word("world", 0.4, 0.8, 0.88),
|
|
]
|
|
seg = self._make_segment("hello world", 0.0, 1.0, words=words)
|
|
result = client._extract_words(seg, 10.0)
|
|
self.assertEqual(len(result), 2)
|
|
self.assertEqual(result[0]["word"], "hello")
|
|
self.assertEqual(result[0]["start"], "10.000")
|
|
self.assertEqual(result[0]["end"], "10.300")
|
|
self.assertEqual(result[0]["probability"], 0.95)
|
|
self.assertEqual(result[1]["word"], "world")
|
|
self.assertEqual(result[1]["start"], "10.400")
|
|
|
|
def test_extract_words_no_words_on_segment(self):
|
|
client = self._make_client(word_timestamps=True)
|
|
seg = self._make_segment("hello", 0.0, 1.0, words=None)
|
|
result = client._extract_words(seg, 0.0)
|
|
self.assertIsNone(result)
|
|
|
|
def test_format_segment_without_words(self):
|
|
client = self._make_client()
|
|
seg = client.format_segment(0.0, 1.0, "hello")
|
|
self.assertNotIn("words", seg)
|
|
|
|
def test_format_segment_with_words(self):
|
|
client = self._make_client(word_timestamps=True)
|
|
words = [{"word": "hello", "start": "0.000", "end": "0.500", "probability": 0.95}]
|
|
seg = client.format_segment(0.0, 1.0, "hello", words=words)
|
|
self.assertIn("words", seg)
|
|
self.assertEqual(len(seg["words"]), 1)
|
|
self.assertEqual(seg["words"][0]["word"], "hello")
|
|
|
|
def test_update_segments_includes_words(self):
|
|
client = self._make_client(word_timestamps=True)
|
|
words1 = [self._make_word("hello", 0.0, 0.5, 0.9)]
|
|
words2 = [self._make_word("world", 0.6, 1.0, 0.85)]
|
|
segments = [
|
|
self._make_segment(" hello", 0.0, 0.5, words=words1),
|
|
self._make_segment(" world", 0.6, 1.0, words=words2),
|
|
]
|
|
last = client.update_segments(segments, 2.0)
|
|
# First segment should be completed (in transcript) with words
|
|
self.assertTrue(len(client.transcript) > 0)
|
|
self.assertIn("words", client.transcript[-1])
|
|
# Last segment should be in-progress with words
|
|
self.assertIsNotNone(last)
|
|
self.assertIn("words", last)
|
|
|
|
def test_update_segments_no_words_when_disabled(self):
|
|
client = self._make_client(word_timestamps=False)
|
|
words1 = [self._make_word("hello", 0.0, 0.5, 0.9)]
|
|
words2 = [self._make_word("world", 0.6, 1.0, 0.85)]
|
|
segments = [
|
|
self._make_segment(" hello", 0.0, 0.5, words=words1),
|
|
self._make_segment(" world", 0.6, 1.0, words=words2),
|
|
]
|
|
last = client.update_segments(segments, 2.0)
|
|
self.assertTrue(len(client.transcript) > 0)
|
|
self.assertNotIn("words", client.transcript[-1])
|
|
self.assertNotIn("words", last)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|