mkbk
218 lines
8.3 KiB
Python
218 lines
8.3 KiB
Python
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
|
|
from owner_voice_pet.models import AudioFrame, AudioSegment, ErrorCode
|
|
from owner_voice_pet.transport import (
|
|
AudioRingBuffer,
|
|
FileReplayTransport,
|
|
MemoryAudioTransport,
|
|
SoundDeviceAudioTransport,
|
|
)
|
|
|
|
|
|
def frame(idx: int, timestamp_ms: int, metadata: dict[str, object] | None = None) -> AudioFrame:
|
|
return AudioFrame(
|
|
pcm=idx.to_bytes(2, "little", signed=False),
|
|
sample_rate=16000,
|
|
channels=1,
|
|
timestamp_ms=timestamp_ms,
|
|
frame_id=idx,
|
|
metadata=metadata or {},
|
|
)
|
|
|
|
|
|
class TransportTests(unittest.TestCase):
|
|
def test_memory_transport_replays_frames_and_captures_playback(self) -> None:
|
|
transport = MemoryAudioTransport([frame(1, 0), frame(2, 20)])
|
|
transport.start_input()
|
|
self.assertEqual([f.frame_id for f in transport.read_frames(10)], [1])
|
|
self.assertEqual([f.frame_id for f in transport.read_frames(10)], [2])
|
|
self.assertEqual(transport.read_frames(10), [])
|
|
|
|
result = transport.play_pcm(AudioSegment(b"\x00\x00", 16000, 1, 0, 100))
|
|
self.assertTrue(result.played)
|
|
self.assertEqual(len(transport.played_segments), 1)
|
|
|
|
def test_memory_transport_reports_missing_output(self) -> None:
|
|
transport = MemoryAudioTransport(output_available=False)
|
|
result = transport.play_pcm(AudioSegment(b"\x00\x00", 16000, 1, 0, 100))
|
|
self.assertFalse(result.played)
|
|
self.assertEqual(result.error.code, ErrorCode.AUDIO_OUTPUT_DEVICE_MISSING)
|
|
|
|
def test_memory_transport_flushes_input_frames(self) -> None:
|
|
transport = MemoryAudioTransport([frame(1, 0), frame(2, 20)])
|
|
transport.start_input()
|
|
self.assertEqual(transport.flush_input(), 2)
|
|
self.assertEqual(transport.flush_count, 1)
|
|
self.assertEqual(transport.read_frames(10), [])
|
|
|
|
def test_memory_transport_can_preserve_prefilled_frames_on_flush(self) -> None:
|
|
transport = MemoryAudioTransport([frame(1, 0), frame(2, 20)], flush_clears_input=False)
|
|
transport.start_input()
|
|
self.assertEqual(transport.flush_input(), 0)
|
|
self.assertEqual(transport.flush_count, 1)
|
|
self.assertEqual([f.frame_id for f in transport.read_frames(10)], [1])
|
|
|
|
def test_file_replay_roundtrip(self) -> None:
|
|
frames = [frame(1, 0, {"wake": True}), frame(2, 20, {"speech": True})]
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
path = Path(tmp) / "fixture.jsonl"
|
|
FileReplayTransport.write_jsonl(path, frames)
|
|
transport = FileReplayTransport.from_jsonl(path)
|
|
transport.start_input()
|
|
self.assertTrue(transport.read_frames(10)[0].metadata["wake"])
|
|
self.assertTrue(transport.read_frames(10)[0].metadata["speech"])
|
|
|
|
def test_ring_buffer_trims_old_frames_and_requires_order(self) -> None:
|
|
buffer = AudioRingBuffer(max_duration_ms=40)
|
|
buffer.extend([frame(1, 0), frame(2, 20), frame(3, 50)])
|
|
self.assertEqual([f.frame_id for f in buffer.frames()], [2, 3])
|
|
with self.assertRaises(ValueError):
|
|
buffer.append(frame(4, 10))
|
|
|
|
def test_sounddevice_transport_is_safe_without_optional_dependency(self) -> None:
|
|
transport = SoundDeviceAudioTransport()
|
|
health = transport.health()
|
|
self.assertIsInstance(health.input_available, bool)
|
|
self.assertIsInstance(health.output_available, bool)
|
|
|
|
def test_sounddevice_transport_reads_from_raw_input_stream(self) -> None:
|
|
fake = FakeSoundDevice()
|
|
transport = SoundDeviceAudioTransport(sounddevice_module=fake)
|
|
transport.start_input(sample_rate=16000, channels=1)
|
|
frames = transport.read_frames(10)
|
|
transport.stop()
|
|
|
|
self.assertEqual(len(frames), 1)
|
|
self.assertEqual(frames[0].pcm, b"\x01\x00\x02\x00")
|
|
self.assertEqual(frames[0].sample_rate, 16000)
|
|
|
|
def test_sounddevice_transport_returns_queued_frame_batch(self) -> None:
|
|
fake = FakeSoundDevice(callback_payloads=[b"\x01\x00", b"\x02\x00", b"\x03\x00"])
|
|
transport = SoundDeviceAudioTransport(sounddevice_module=fake)
|
|
transport.start_input(sample_rate=16000, channels=1)
|
|
frames = transport.read_frames(10)
|
|
transport.stop()
|
|
|
|
self.assertEqual([item.frame_id for item in frames], [0, 1, 2])
|
|
|
|
def test_sounddevice_transport_flushes_queued_input(self) -> None:
|
|
fake = FakeSoundDevice()
|
|
transport = SoundDeviceAudioTransport(sounddevice_module=fake)
|
|
transport.start_input(sample_rate=16000, channels=1)
|
|
self.assertEqual(transport.flush_input(), 1)
|
|
self.assertEqual(transport.read_frames(0), [])
|
|
transport.stop()
|
|
|
|
def test_sounddevice_transport_plays_pcm_to_raw_output_stream(self) -> None:
|
|
fake = FakeSoundDevice()
|
|
transport = SoundDeviceAudioTransport(sounddevice_module=fake)
|
|
result = transport.play_pcm(AudioSegment(b"\x00\x00\x01\x00", 16000, 1, 0, 20))
|
|
|
|
self.assertTrue(result.played)
|
|
self.assertEqual(fake.output_writes, [b"\x00\x00\x01\x00"])
|
|
self.assertEqual(fake.output_stream_open_count, 1)
|
|
|
|
def test_sounddevice_chunk_playback_keeps_one_output_stream_open(self) -> None:
|
|
fake = FakeSoundDevice()
|
|
transport = SoundDeviceAudioTransport(sounddevice_module=fake)
|
|
segment = AudioSegment(b"\x00\x00" * 1600, 16000, 1, 0, 100)
|
|
callbacks: list[int] = []
|
|
|
|
result = transport.play_pcm_chunks(
|
|
segment,
|
|
chunk_ms=20,
|
|
after_chunk=lambda _chunk, elapsed_ms: callbacks.append(elapsed_ms) is not None and False,
|
|
)
|
|
|
|
self.assertTrue(result.played)
|
|
self.assertEqual(fake.output_stream_open_count, 1)
|
|
self.assertGreater(len(fake.output_writes), 1)
|
|
self.assertEqual(callbacks[-1], result.duration_ms)
|
|
|
|
def test_sounddevice_chunk_playback_stops_on_signal(self) -> None:
|
|
fake = FakeSoundDevice()
|
|
transport = SoundDeviceAudioTransport(sounddevice_module=fake)
|
|
segment = AudioSegment(b"\x00\x00" * 1600, 16000, 1, 0, 100)
|
|
|
|
result = transport.play_pcm_chunks(
|
|
segment,
|
|
chunk_ms=20,
|
|
should_stop=lambda: len(fake.output_writes) >= 2,
|
|
)
|
|
|
|
self.assertTrue(result.played)
|
|
self.assertEqual(len(fake.output_writes), 2)
|
|
self.assertEqual(result.duration_ms, 40)
|
|
|
|
def test_sounddevice_device_report_uses_query_devices(self) -> None:
|
|
fake = FakeSoundDevice()
|
|
report = SoundDeviceAudioTransport(sounddevice_module=fake).device_report()
|
|
|
|
self.assertTrue(report["ok"])
|
|
self.assertEqual(len(report["devices"]), 2)
|
|
|
|
|
|
class FakeInputStream:
|
|
def __init__(self, callback, **kwargs) -> None:
|
|
self.callback = callback
|
|
self.kwargs = kwargs
|
|
self.started = False
|
|
self.closed = False
|
|
|
|
def start(self) -> None:
|
|
self.started = True
|
|
payloads = self.kwargs.get("callback_payloads") or [b"\x01\x00\x02\x00"]
|
|
for payload in payloads:
|
|
self.callback(payload, max(1, len(payload) // 2), None, "")
|
|
|
|
def stop(self) -> None:
|
|
self.started = False
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
class FakeOutputStream:
|
|
def __init__(self, owner, **kwargs) -> None:
|
|
self.owner = owner
|
|
self.kwargs = kwargs
|
|
|
|
def __enter__(self):
|
|
return self
|
|
|
|
def __exit__(self, *args) -> None:
|
|
return None
|
|
|
|
def write(self, data: bytes) -> None:
|
|
self.owner.output_writes.append(bytes(data))
|
|
|
|
|
|
class FakeSoundDevice:
|
|
def __init__(self, callback_payloads: list[bytes] | None = None) -> None:
|
|
self.callback_payloads = callback_payloads
|
|
self.output_writes: list[bytes] = []
|
|
self.output_stream_open_count = 0
|
|
|
|
def RawInputStream(self, **kwargs):
|
|
if self.callback_payloads is not None:
|
|
kwargs["callback_payloads"] = self.callback_payloads
|
|
return FakeInputStream(**kwargs)
|
|
|
|
def RawOutputStream(self, **kwargs):
|
|
self.output_stream_open_count += 1
|
|
return FakeOutputStream(self, **kwargs)
|
|
|
|
def query_devices(self):
|
|
return [
|
|
{"name": "Mic", "max_input_channels": 1, "max_output_channels": 0, "default_samplerate": 16000},
|
|
{"name": "Speaker", "max_input_channels": 0, "max_output_channels": 2, "default_samplerate": 48000},
|
|
]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|