fix unit tests
This commit is contained in:
@@ -1,23 +1,25 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import os
|
import unittest
|
||||||
import wave
|
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
try:
|
||||||
|
import pyrnnoise
|
||||||
|
except ImportError:
|
||||||
|
pyrnnoise = None
|
||||||
|
|
||||||
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
|
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
|
||||||
from pipecat.frames.frames import FilterEnableFrame
|
from pipecat.frames.frames import FilterEnableFrame
|
||||||
|
|
||||||
|
|
||||||
async def test_rnnoise_cancellation_functionality():
|
class TestRNNoiseCancellation(unittest.IsolatedAsyncioTestCase):
|
||||||
|
async def test_rnnoise_cancellation_functionality(self):
|
||||||
print("\nStarting Noise Cancellation Test")
|
print("\nStarting Noise Cancellation Test")
|
||||||
|
|
||||||
# 1. Check for pyrnnoise
|
# 1. Check for pyrnnoise
|
||||||
try:
|
if pyrnnoise is None:
|
||||||
import pyrnnoise
|
self.skipTest("pyrnnoise not installed. Cannot verify actual noise cancellation.")
|
||||||
except ImportError:
|
|
||||||
pytest.skip("pyrnnoise not installed. Cannot verify actual noise cancellation.")
|
|
||||||
return
|
|
||||||
|
|
||||||
# 2. Generate clean speech-like audio (Harmonic series)
|
# 2. Generate clean speech-like audio (Harmonic series)
|
||||||
sample_rate = 48000
|
sample_rate = 48000
|
||||||
@@ -140,7 +142,9 @@ async def test_rnnoise_cancellation_functionality():
|
|||||||
noise_power_output = np.mean(aligned_output[silent_mask].astype(float) ** 2)
|
noise_power_output = np.mean(aligned_output[silent_mask].astype(float) ** 2)
|
||||||
print(f"Noise Power in Silence (Input): {noise_power_input:.2f}")
|
print(f"Noise Power in Silence (Input): {noise_power_input:.2f}")
|
||||||
print(f"Noise Power in Silence (Output): {noise_power_output:.2f}")
|
print(f"Noise Power in Silence (Output): {noise_power_output:.2f}")
|
||||||
assert noise_power_output < noise_power_input, "Noise power in silence not reduced"
|
self.assertLess(
|
||||||
|
noise_power_output, noise_power_input, "Noise power in silence not reduced"
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
print("Warning: Not enough silent samples found for noise floor check.")
|
print("Warning: Not enough silent samples found for noise floor check.")
|
||||||
|
|
||||||
@@ -154,14 +158,12 @@ async def test_rnnoise_cancellation_functionality():
|
|||||||
"Warning: Overall MSE did not improve (speech distortion?). Relying on Noise Power check."
|
"Warning: Overall MSE did not improve (speech distortion?). Relying on Noise Power check."
|
||||||
)
|
)
|
||||||
# If we passed the noise power check above, we are good.
|
# If we passed the noise power check above, we are good.
|
||||||
assert np.sum(silent_mask) > 1000 and np.mean(
|
self.assertTrue(
|
||||||
aligned_output[silent_mask].astype(float) ** 2
|
np.sum(silent_mask) > 1000
|
||||||
) < np.mean(aligned_noisy[silent_mask].astype(float) ** 2)
|
and np.mean(aligned_output[silent_mask].astype(float) ** 2)
|
||||||
|
< np.mean(aligned_noisy[silent_mask].astype(float) ** 2)
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
assert mse_output < mse_input, "MSE did not improve"
|
self.assertLess(mse_output, mse_input, "MSE did not improve")
|
||||||
|
|
||||||
print("Test Passed: Noise cancellation verified.")
|
print("Test Passed: Noise cancellation verified.")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
asyncio.run(test_rnnoise_cancellation_functionality())
|
|
||||||
|
|||||||
@@ -8,11 +8,20 @@ import unittest
|
|||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
|
|
||||||
|
try:
|
||||||
|
import pyrnnoise
|
||||||
|
except ImportError:
|
||||||
|
pyrnnoise = None
|
||||||
|
|
||||||
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
|
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
|
||||||
from pipecat.frames.frames import FilterEnableFrame
|
from pipecat.frames.frames import FilterEnableFrame
|
||||||
|
|
||||||
|
|
||||||
class TestRNNoiseFilter(unittest.IsolatedAsyncioTestCase):
|
class TestRNNoiseFilter(unittest.IsolatedAsyncioTestCase):
|
||||||
|
def setUp(self):
|
||||||
|
if pyrnnoise is None:
|
||||||
|
self.skipTest("pyrnnoise not installed")
|
||||||
|
|
||||||
async def test_rnnoise_filter_reduces_noise(self):
|
async def test_rnnoise_filter_reduces_noise(self):
|
||||||
"""Test that RNNoise filter reduces noise in audio."""
|
"""Test that RNNoise filter reduces noise in audio."""
|
||||||
filter = RNNoiseFilter()
|
filter = RNNoiseFilter()
|
||||||
|
|||||||
@@ -1,26 +1,24 @@
|
|||||||
import asyncio
|
import asyncio
|
||||||
import sys
|
import sys
|
||||||
from unittest.mock import MagicMock
|
import unittest
|
||||||
|
from unittest.mock import MagicMock, patch
|
||||||
|
|
||||||
import numpy as np
|
import numpy as np
|
||||||
import pytest
|
|
||||||
|
|
||||||
# Mock pyrnnoise BEFORE importing RNNoiseFilter
|
# We don't need to mock sys.modules here if we use patch on the imported module member
|
||||||
mock_pyrnnoise = MagicMock()
|
# But we need to ensure RNNoiseFilter is imported so we can patch its member
|
||||||
mock_rnnoise_class = MagicMock()
|
|
||||||
mock_pyrnnoise.RNNoise = mock_rnnoise_class
|
|
||||||
sys.modules["pyrnnoise"] = mock_pyrnnoise
|
|
||||||
|
|
||||||
# Now import the filter
|
|
||||||
try:
|
try:
|
||||||
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
|
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
|
||||||
from pipecat.frames.frames import FilterEnableFrame
|
from pipecat.frames.frames import FilterEnableFrame
|
||||||
except ImportError as e:
|
except ImportError as e:
|
||||||
|
# If dependencies are missing (like numpy?), we can't test
|
||||||
print(f"Failed to import RNNoiseFilter: {e}")
|
print(f"Failed to import RNNoiseFilter: {e}")
|
||||||
sys.exit(1)
|
sys.exit(1)
|
||||||
|
|
||||||
|
|
||||||
async def test_rnnoise_resampling_16k_to_48k_and_back():
|
class TestRNNoiseResampling(unittest.IsolatedAsyncioTestCase):
|
||||||
|
@patch("pipecat.audio.filters.rnnoise_filter.RNNoise")
|
||||||
|
async def test_rnnoise_resampling_16k_to_48k_and_back(self, mock_rnnoise_class):
|
||||||
print("\nStarting Resampling Test: 16kHz -> 48kHz -> 16kHz")
|
print("\nStarting Resampling Test: 16kHz -> 48kHz -> 16kHz")
|
||||||
|
|
||||||
# Configure Mock with buffering behavior
|
# Configure Mock with buffering behavior
|
||||||
@@ -56,9 +54,12 @@ async def test_rnnoise_resampling_16k_to_48k_and_back():
|
|||||||
audio_data = (np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16)
|
audio_data = (np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16)
|
||||||
audio_bytes = audio_data.tobytes()
|
audio_bytes = audio_data.tobytes()
|
||||||
|
|
||||||
print(f"Input audio: {len(audio_bytes)} bytes, {len(audio_data)} samples at {sample_rate}Hz")
|
print(
|
||||||
|
f"Input audio: {len(audio_bytes)} bytes, {len(audio_data)} samples at {sample_rate}Hz"
|
||||||
|
)
|
||||||
|
|
||||||
# 2. Initialize RNNoiseFilter
|
# 2. Initialize RNNoiseFilter
|
||||||
|
# This will use the patched RNNoise
|
||||||
rnnoise_filter = RNNoiseFilter()
|
rnnoise_filter = RNNoiseFilter()
|
||||||
await rnnoise_filter.start(sample_rate)
|
await rnnoise_filter.start(sample_rate)
|
||||||
|
|
||||||
@@ -81,38 +82,32 @@ async def test_rnnoise_resampling_16k_to_48k_and_back():
|
|||||||
|
|
||||||
# 4. Verify output length
|
# 4. Verify output length
|
||||||
# Expect roughly same length
|
# Expect roughly same length
|
||||||
# Input: 16000 samples.
|
|
||||||
# Upsampled to 48000.
|
|
||||||
# 48000 / 480 = 100 chunks.
|
|
||||||
# So we expect roughly 100 calls to process_chunk.
|
|
||||||
expected_chunks = (len(audio_data) * 48000 // sample_rate) // 480
|
expected_chunks = (len(audio_data) * 48000 // sample_rate) // 480
|
||||||
print(f"Expected chunks: ~{expected_chunks}")
|
print(f"Expected chunks: ~{expected_chunks}")
|
||||||
|
|
||||||
# Check that we actually processed something
|
# Check that we actually processed something
|
||||||
assert processed_chunks_count >= expected_chunks - 5, "Too few chunks processed"
|
self.assertGreaterEqual(
|
||||||
|
processed_chunks_count, expected_chunks - 5, "Too few chunks processed"
|
||||||
|
)
|
||||||
|
|
||||||
# Check output length
|
# Check output length
|
||||||
assert len(processed_audio) > 0, "Output should not be empty"
|
self.assertGreater(len(processed_audio), 0, "Output should not be empty")
|
||||||
|
|
||||||
# Check length matches input (with some tolerance for buffering latency)
|
# Check length matches input (with some tolerance for buffering latency)
|
||||||
# Since we don't flush the filter explicitly (no flush method in RNNoiseFilter yet),
|
|
||||||
# some data might remain in buffers.
|
|
||||||
# Max loss:
|
|
||||||
# - Resampler input buffer
|
|
||||||
# - RNNoise buffer (max 480 samples = 10ms)
|
|
||||||
# - Resampler output buffer
|
|
||||||
|
|
||||||
# 100ms tolerance?
|
# 100ms tolerance?
|
||||||
byte_tolerance = int(0.2 * sample_rate * 2)
|
byte_tolerance = int(0.2 * sample_rate * 2)
|
||||||
assert len(processed_audio) >= len(audio_bytes) - byte_tolerance, (
|
self.assertGreaterEqual(
|
||||||
f"Output too short: {len(processed_audio)} vs {len(audio_bytes)}"
|
len(processed_audio),
|
||||||
|
len(audio_bytes) - byte_tolerance,
|
||||||
|
f"Output too short: {len(processed_audio)} vs {len(audio_bytes)}",
|
||||||
)
|
)
|
||||||
assert len(processed_audio) <= len(audio_bytes) + byte_tolerance, (
|
self.assertLessEqual(
|
||||||
f"Output too long: {len(processed_audio)} vs {len(audio_bytes)}"
|
len(processed_audio),
|
||||||
|
len(audio_bytes) + byte_tolerance,
|
||||||
|
f"Output too long: {len(processed_audio)} vs {len(audio_bytes)}",
|
||||||
)
|
)
|
||||||
|
|
||||||
# 5. Check sample rate / pitch preservation
|
# 5. Check sample rate / pitch preservation
|
||||||
# If we upsampled and downsampled correctly, the pitch should be 440Hz.
|
|
||||||
output_data = np.frombuffer(processed_audio, dtype=np.int16)
|
output_data = np.frombuffer(processed_audio, dtype=np.int16)
|
||||||
|
|
||||||
if len(output_data) > 2000:
|
if len(output_data) > 2000:
|
||||||
@@ -127,10 +122,8 @@ async def test_rnnoise_resampling_16k_to_48k_and_back():
|
|||||||
peak_freq = freqs[peak_idx]
|
peak_freq = freqs[peak_idx]
|
||||||
|
|
||||||
print(f"Peak frequency: {peak_freq:.2f} Hz")
|
print(f"Peak frequency: {peak_freq:.2f} Hz")
|
||||||
assert abs(peak_freq - 440) < 50, f"Frequency shifted significantly: {peak_freq} vs 440"
|
self.assertLess(
|
||||||
|
abs(peak_freq - 440), 50, f"Frequency shifted significantly: {peak_freq} vs 440"
|
||||||
|
)
|
||||||
|
|
||||||
print("Test Passed: Resampling logic verified (with mocked RNNoise).")
|
print("Test Passed: Resampling logic verified (with mocked RNNoise).")
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
asyncio.run(test_rnnoise_resampling_16k_to_48k_and_back())
|
|
||||||
|
|||||||
Reference in New Issue
Block a user