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()