345 lines
16 KiB
Python
345 lines
16 KiB
Python
import json
|
|
import tempfile
|
|
import threading
|
|
import time
|
|
import unittest
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from unittest.mock import patch
|
|
|
|
import numpy as np
|
|
|
|
import blurt_session
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Config:
|
|
model: str = "test-model"
|
|
language: str = "en"
|
|
api_key: str = "do-not-store"
|
|
|
|
|
|
@dataclass
|
|
class Chunk:
|
|
samples: np.ndarray
|
|
start: int
|
|
end: int
|
|
speech: bool = True
|
|
|
|
|
|
class FakeAudio:
|
|
def __init__(self, rate, channels):
|
|
self.position = 0
|
|
self.tail = np.empty(0, dtype=np.int16)
|
|
|
|
def feed(self, frames):
|
|
value = int(np.asarray(frames).flat[0])
|
|
samples = np.full(100, value, dtype=np.int16)
|
|
start = self.position
|
|
self.position += len(samples)
|
|
return [Chunk(samples, start, self.position, value != 0)]
|
|
|
|
def finish(self):
|
|
return []
|
|
|
|
def checkpoint(self):
|
|
return self.tail.copy()
|
|
|
|
|
|
class SessionTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
|
|
def scheduler(self, request, workers=3):
|
|
scheduler = blurt_session.TranscriptionScheduler(request, workers=workers)
|
|
self.addCleanup(scheduler.shutdown)
|
|
return scheduler
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_transcribes_before_stop_and_orders_out_of_order_results(self):
|
|
first_started = threading.Event()
|
|
release_first = threading.Event()
|
|
|
|
def request(audio, config):
|
|
value = int(audio[0])
|
|
if value == 1:
|
|
first_started.set()
|
|
release_first.wait(2)
|
|
return f"part {value}"
|
|
|
|
completed = []
|
|
scheduler = self.scheduler(request)
|
|
session = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name,
|
|
scheduler, on_complete=completed.append)
|
|
self.assertTrue(session.feed(np.array([[1]], dtype=np.float32)))
|
|
self.assertTrue(first_started.wait(1), "upload did not begin while recording")
|
|
self.assertTrue(session.feed(np.array([[2]], dtype=np.float32)))
|
|
deadline = time.monotonic() + 1
|
|
while session.done < 1 and time.monotonic() < deadline:
|
|
time.sleep(0.005)
|
|
self.assertEqual(session.state, "recording")
|
|
self.assertEqual(completed, [], "must not deliver during recording")
|
|
session.stop()
|
|
release_first.set()
|
|
result = session.join(2)
|
|
self.assertIsNotNone(result)
|
|
self.assertEqual(result.text, "part 1 part 2")
|
|
self.assertTrue(result.complete)
|
|
self.assertEqual(completed, [session])
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_failure_is_recoverable_and_partial_ack_keeps_spool(self):
|
|
scheduler = self.scheduler(lambda audio, config: (_ for _ in ()).throw(RuntimeError("offline")))
|
|
session = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name, scheduler)
|
|
session.feed(np.array([[7]], dtype=np.float32))
|
|
session.stop()
|
|
result = session.join(2)
|
|
self.assertFalse(result.complete)
|
|
self.assertIn("offline", result.errors[0])
|
|
self.assertEqual(result.text, "")
|
|
session.acknowledge_delivery()
|
|
self.assertTrue(session.path.exists())
|
|
manifest = json.loads((session.path / "session.json").read_text())
|
|
self.assertNotIn("api_key", manifest["config"])
|
|
self.assertEqual(manifest["config"]["model"], "test-model")
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_recovery_retries_failed_audio_and_preserves_successful_text(self):
|
|
calls = []
|
|
|
|
def initial(audio, config):
|
|
value = int(audio[0])
|
|
if value == 2:
|
|
raise OSError("temporary")
|
|
return "one"
|
|
|
|
first_scheduler = blurt_session.TranscriptionScheduler(initial, workers=2)
|
|
session = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name, first_scheduler)
|
|
session.feed(np.array([[1]], dtype=np.float32))
|
|
session.feed(np.array([[2]], dtype=np.float32))
|
|
session.stop()
|
|
self.assertFalse(session.join(2).complete)
|
|
session.release_for_retry()
|
|
first_scheduler.shutdown()
|
|
|
|
def recovered(audio, config):
|
|
calls.append(int(audio[0]))
|
|
self.assertEqual(config.language, "en")
|
|
self.assertEqual(config.api_key, "new-key")
|
|
return "two"
|
|
|
|
second_scheduler = self.scheduler(recovered)
|
|
records = blurt_session.RecoveryStore(self.temp.name).discover()
|
|
self.assertEqual(len(records), 1)
|
|
resumed = blurt_session.RecoveryStore(self.temp.name).resume(
|
|
records[0], Config(api_key="new-key"), second_scheduler)
|
|
result = resumed.join(2)
|
|
self.assertTrue(result.complete)
|
|
self.assertEqual(result.text, "one two")
|
|
self.assertEqual(calls, [2])
|
|
resumed.acknowledge_delivery()
|
|
self.assertFalse(resumed.path.exists())
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_silence_is_spooled_but_not_uploaded(self):
|
|
called = []
|
|
session = blurt_session.RecordingSession(
|
|
16000, 1, Config(), self.temp.name,
|
|
self.scheduler(lambda audio, config: called.append(True) or "bad"))
|
|
session.feed(np.array([[0]], dtype=np.float32))
|
|
session.stop()
|
|
result = session.join(2)
|
|
self.assertTrue(result.complete)
|
|
self.assertEqual(result.text, "")
|
|
self.assertEqual(called, [])
|
|
self.assertTrue(any(session.path.glob("chunk-*.wav")))
|
|
|
|
def test_discovery_includes_checkpoint_tail(self):
|
|
path = Path(self.temp.name) / "crashed"
|
|
path.mkdir()
|
|
blurt_session._write_wav(path / "tail.wav", np.full(20, 2, dtype=np.int16))
|
|
(path / "session.json").write_text(json.dumps({
|
|
"version": 1, "session_id": "crashed", "state": "recording",
|
|
"config": {"model": "old", "language": "de"}, "done": 0,
|
|
"total": 0, "entries": [], "has_tail": True,
|
|
}))
|
|
scheduler = self.scheduler(lambda audio, config: "tail")
|
|
record = blurt_session.RecoveryStore(self.temp.name).discover()[0]
|
|
self.assertEqual(record.config["language"], "de")
|
|
resumed = blurt_session.RecordingSession.recover(path, Config(api_key="fresh"), scheduler)
|
|
self.assertEqual(resumed.join(2).text, "tail")
|
|
|
|
def _interrupted(self, name="crash", entries=None, **extra):
|
|
path = Path(self.temp.name) / name
|
|
path.mkdir()
|
|
data = {"version": 1, "session_id": name, "state": "recording",
|
|
"config": {"model": "old", "language": "de"}, "done": 0,
|
|
"total": 0, "entries": entries or [], "has_tail": False,
|
|
"tail_start": 0, "created_at": extra.pop("created_at", time.time())}
|
|
data.update(extra)
|
|
(path / "session.json").write_text(json.dumps(data))
|
|
return path
|
|
|
|
def test_recovery_adopts_orphan_sealed_wav_instead_of_old_tail(self):
|
|
path = self._interrupted()
|
|
blurt_session._write_wav(path / "chunk-000000.wav", np.full(20, 4, np.int16))
|
|
blurt_session._write_wav(path / "tail.wav", np.full(10, 2, np.int16))
|
|
seen = []
|
|
scheduler = self.scheduler(lambda audio, config: seen.append((len(audio), int(audio[0]))) or "sealed")
|
|
result = blurt_session.RecordingSession.recover(path, Config(), scheduler).join(2)
|
|
self.assertEqual(result.text, "sealed")
|
|
self.assertEqual(seen, [(20, 4)])
|
|
|
|
def test_silent_interrupted_tail_is_not_uploaded(self):
|
|
path = self._interrupted()
|
|
blurt_session._write_wav(path / "tail.wav", np.zeros(20, np.int16))
|
|
calls = []
|
|
session = blurt_session.RecordingSession.recover(
|
|
path, Config(), self.scheduler(lambda audio, config: calls.append(1) or "wrong"))
|
|
self.assertTrue(session.join(2).complete)
|
|
self.assertEqual(session.result.text, "")
|
|
self.assertEqual(calls, [])
|
|
|
|
def test_missing_recovery_file_finishes_with_error(self):
|
|
entry = {"index": 0, "file": "gone.wav", "speech": True, "status": "pending"}
|
|
path = self._interrupted(entries=[entry], total=1)
|
|
session = blurt_session.RecordingSession.recover(path, Config(), self.scheduler(lambda *_: "wrong"))
|
|
result = session.join(2)
|
|
self.assertFalse(result.complete)
|
|
self.assertEqual(result.text, "")
|
|
self.assertIn("missing", result.errors[0])
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_raising_callbacks_do_not_stop_worker_pool(self):
|
|
scheduler = self.scheduler(lambda audio, config: str(int(audio[0])), workers=1)
|
|
bad = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name, scheduler,
|
|
on_update=lambda _: (_ for _ in ()).throw(RuntimeError("ui")),
|
|
on_complete=lambda _: (_ for _ in ()).throw(RuntimeError("ui")))
|
|
bad.feed(np.array([[3]], np.float32)); bad.stop()
|
|
self.assertEqual(bad.join(2).text, "3")
|
|
good = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name, scheduler)
|
|
good.feed(np.array([[4]], np.float32)); good.stop()
|
|
self.assertEqual(good.join(2).text, "4")
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_active_exclusion_and_release_for_retry(self):
|
|
session = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name,
|
|
self.scheduler(lambda *_: "text"))
|
|
self.assertEqual(blurt_session.RecoveryStore(self.temp.name).discover(), [])
|
|
session.feed(np.array([[1]], np.float32)); session.stop(); session.join(2)
|
|
self.assertEqual(blurt_session.RecoveryStore(self.temp.name).discover(), [])
|
|
session.release_for_retry()
|
|
self.assertEqual(len(blurt_session.RecoveryStore(self.temp.name).discover()), 1)
|
|
|
|
def test_from_wav_migrates_source_after_durable_spool(self):
|
|
source = Path(self.temp.name) / "legacy.wav"
|
|
blurt_session._write_wav(source, np.tile(np.array([-1000, 1000], np.int16), 800))
|
|
spool = Path(self.temp.name) / "sessions"; spool.mkdir()
|
|
session = blurt_session.RecordingSession.from_wav(
|
|
source, Config(), spool, self.scheduler(lambda *_: "legacy"))
|
|
self.assertFalse(source.exists())
|
|
self.assertTrue((session.path / "session.json").exists())
|
|
self.assertEqual(session.join(2).text, "legacy")
|
|
|
|
def test_discovery_is_oldest_first(self):
|
|
self._interrupted("new", created_at=20)
|
|
self._interrupted("old", created_at=10)
|
|
self.assertEqual([r.session_id for r in blurt_session.RecoveryStore(self.temp.name).discover()],
|
|
["old", "new"])
|
|
|
|
def test_old_tail_overlap_keeps_only_suffix_after_orphan(self):
|
|
path = self._interrupted(tail_start=0)
|
|
blurt_session._write_wav(path / "chunk-000000.wav", np.full(6, 3, np.int16))
|
|
blurt_session._write_wav(path / "tail.wav", np.arange(1, 11, dtype=np.int16))
|
|
seen = []
|
|
scheduler = self.scheduler(lambda audio, config: seen.append(audio.copy()) or "x")
|
|
session = blurt_session.RecordingSession.recover(path, Config(), scheduler)
|
|
self.assertEqual(session.join(2).text, "x x")
|
|
self.assertCountEqual([tuple(int(v) for v in x) for x in seen],
|
|
[(3, 3, 3, 3, 3, 3), (7, 8, 9, 10)])
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_wav_persistence_failure_finishes_without_waiting_forever(self):
|
|
scheduler = self.scheduler(lambda *_: "unused")
|
|
session = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name, scheduler)
|
|
with patch.object(blurt_session, "_write_wav", side_effect=OSError("disk full")):
|
|
session.feed(np.array([[5]], np.float32)); session.stop()
|
|
result = session.join(2)
|
|
self.assertIsNotNone(result)
|
|
self.assertFalse(result.complete)
|
|
self.assertIn("disk full", result.errors[0])
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_inflight_job_is_not_failed_by_later_capture_error(self):
|
|
started = threading.Event(); release = threading.Event()
|
|
def request(audio, config):
|
|
started.set(); release.wait(2); return "first"
|
|
scheduler = self.scheduler(request, workers=1)
|
|
session = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name, scheduler)
|
|
real_write = blurt_session._write_wav
|
|
writes = 0
|
|
def fail_second(path, samples):
|
|
nonlocal writes
|
|
writes += 1
|
|
if writes == 2:
|
|
raise OSError("disk full")
|
|
return real_write(path, samples)
|
|
with patch.object(blurt_session, "_write_wav", side_effect=fail_second):
|
|
session.feed(np.array([[1]], np.float32)); self.assertTrue(started.wait(1))
|
|
session.feed(np.array([[2]], np.float32)); session.stop()
|
|
self.assertIsNone(session.join(0.05))
|
|
release.set()
|
|
result = session.join(2)
|
|
self.assertEqual(result.text, "first […]")
|
|
self.assertEqual(session.done, session.total)
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_input_status_makes_result_incomplete(self):
|
|
session = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name,
|
|
self.scheduler(lambda *_: "kept"))
|
|
session.report_input_status("input overflow")
|
|
session.feed(np.array([[1]], np.float32)); session.stop()
|
|
result = session.join(2)
|
|
self.assertFalse(result.complete)
|
|
self.assertEqual(result.text, "kept […]")
|
|
|
|
@patch.object(blurt_session, "StreamingAudio", FakeAudio)
|
|
def test_partial_ack_retains_successful_text(self):
|
|
values = iter(["saved", RuntimeError("bad")])
|
|
def request(*_):
|
|
value = next(values)
|
|
if isinstance(value, Exception): raise value
|
|
return value
|
|
session = blurt_session.RecordingSession(16000, 1, Config(), self.temp.name,
|
|
self.scheduler(request, workers=1))
|
|
session.feed(np.array([[1]], np.float32)); session.feed(np.array([[2]], np.float32)); session.stop()
|
|
self.assertFalse(session.join(2).complete)
|
|
session.acknowledge_delivery()
|
|
manifest = json.loads((session.path / "session.json").read_text())
|
|
self.assertEqual(manifest["entries"][0]["text"], "saved")
|
|
|
|
def test_feed_queue_overflow_returns_immediately_and_records_gap(self):
|
|
entered = threading.Event(); release = threading.Event()
|
|
class SlowAudio(FakeAudio):
|
|
def feed(self, frames):
|
|
entered.set(); release.wait(2)
|
|
return super().feed(frames)
|
|
with patch.object(blurt_session, "StreamingAudio", SlowAudio):
|
|
session = blurt_session.RecordingSession(
|
|
16000, 1, Config(), self.temp.name, self.scheduler(lambda *_: "first"),
|
|
input_queue_size=1)
|
|
session.feed(np.array([[1]], np.float32)); self.assertTrue(entered.wait(1))
|
|
self.assertTrue(session.feed(np.array([[2]], np.float32)))
|
|
started = time.monotonic()
|
|
self.assertFalse(session.feed(np.array([[3]], np.float32)))
|
|
self.assertLess(time.monotonic() - started, 0.05)
|
|
release.set()
|
|
result = session.join(2)
|
|
self.assertFalse(result.complete)
|
|
self.assertIn("overflowed", " ".join(result.errors))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|