125 lines
3.5 KiB
Python
125 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
import time
|
|
import types
|
|
|
|
import numpy as np
|
|
|
|
from src.audio_monitor import MvdrAudioMonitor
|
|
from src.mvdr_beamformer import MvdrResult
|
|
|
|
|
|
class FakeOutputStream:
|
|
def __init__(self, writes: list[np.ndarray], **kwargs) -> None:
|
|
self.kwargs = kwargs
|
|
self.writes = writes
|
|
self.aborted = False
|
|
self.closed = False
|
|
|
|
def start(self) -> None:
|
|
pass
|
|
|
|
def write(self, audio: np.ndarray) -> bool:
|
|
self.writes.append(audio.copy())
|
|
return False
|
|
|
|
def abort(self) -> None:
|
|
self.aborted = True
|
|
|
|
def close(self) -> None:
|
|
self.closed = True
|
|
|
|
|
|
def result(
|
|
timestamp: float,
|
|
sample_end: int,
|
|
value: float = 0.01,
|
|
stream_samples: int = 100,
|
|
) -> MvdrResult:
|
|
audio = np.full(100, value, dtype=np.float32)
|
|
audio.setflags(write=False)
|
|
stream_audio = np.full(stream_samples, value, dtype=np.float32)
|
|
stream_audio.setflags(write=False)
|
|
return MvdrResult(
|
|
audio,
|
|
-40.0,
|
|
-30.0,
|
|
45.0,
|
|
timestamp,
|
|
sample_end,
|
|
stream_audio,
|
|
sample_end - stream_samples,
|
|
sample_end,
|
|
)
|
|
|
|
|
|
def test_monitor_plays_only_new_mvdr_tail_with_gain(monkeypatch) -> None:
|
|
writes: list[np.ndarray] = []
|
|
streams: list[FakeOutputStream] = []
|
|
devices = [
|
|
{
|
|
"name": "Fake speakers",
|
|
"max_output_channels": 2,
|
|
"hostapi": 0,
|
|
}
|
|
]
|
|
|
|
def output_stream(**kwargs):
|
|
stream = FakeOutputStream(writes, **kwargs)
|
|
streams.append(stream)
|
|
return stream
|
|
|
|
fake_sounddevice = types.SimpleNamespace(
|
|
default=types.SimpleNamespace(device=(-1, 0)),
|
|
query_devices=lambda: devices,
|
|
query_hostapis=lambda _index: {"name": "Fake API"},
|
|
check_output_settings=lambda **_kwargs: None,
|
|
OutputStream=output_stream,
|
|
)
|
|
monkeypatch.setitem(sys.modules, "sounddevice", fake_sounddevice)
|
|
|
|
current = [result(1.0, 100)]
|
|
monitor = MvdrAudioMonitor(24_000, lambda: current[0])
|
|
monitor.set_gain_db(20.0)
|
|
monitor.set_enabled(True)
|
|
try:
|
|
deadline = time.monotonic() + 1.0
|
|
while len(writes) < 1 and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
assert writes[0].shape == (100, 1)
|
|
np.testing.assert_allclose(writes[0][:, 0], 0.1, atol=1e-6)
|
|
|
|
current[0] = result(2.0, 140, value=0.02, stream_samples=40)
|
|
deadline = time.monotonic() + 1.0
|
|
while len(writes) < 2 and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
assert writes[1].shape == (40, 1)
|
|
np.testing.assert_allclose(writes[1][:, 0], 0.2, atol=1e-6)
|
|
finally:
|
|
monitor.close()
|
|
|
|
assert streams[0].kwargs["device"] == 0
|
|
assert streams[0].kwargs["samplerate"] == 24_000
|
|
assert streams[0].aborted
|
|
assert streams[0].closed
|
|
|
|
|
|
def test_monitor_gain_range_and_clipping(monkeypatch) -> None:
|
|
fake_sounddevice = types.SimpleNamespace(
|
|
default=types.SimpleNamespace(device=(-1, -1)),
|
|
query_devices=lambda: [],
|
|
)
|
|
monkeypatch.setitem(sys.modules, "sounddevice", fake_sounddevice)
|
|
monitor = MvdrAudioMonitor(24_000, lambda: None)
|
|
try:
|
|
assert monitor.set_gain_db(-50.0) == -10.0
|
|
assert monitor.set_gain_db(80.0) == 40.0
|
|
amplified, clipped = monitor.apply_gain(
|
|
np.array([0.02, -0.02], dtype=np.float32), 40.0
|
|
)
|
|
np.testing.assert_array_equal(amplified, np.array([1.0, -1.0]))
|
|
assert clipped
|
|
finally:
|
|
monitor.close()
|