add tests, thx gpt

Author: theomonnom

dev-requirements.txt

@@ -9,3 +9,5 @@ auditwheel; sys_platform == 'linux'
 cibuildwheel
 
 pytest
+
+matplotlib

livekit-rtc/livekit/rtc/audio_mixer.py

@@ -1,12 +1,11 @@
-from typing import AsyncIterator, Optional
 import asyncio
 import numpy as np
 import contextlib
 from dataclasses import dataclass
+from typing import AsyncIterator, Optional
 from .audio_frame import AudioFrame
 from .log import logger
 
-
 _Stream = AsyncIterator[AudioFrame]
 
 
@@ -58,7 +57,9 @@ def __init__(
         self._chunk_size: int = blocksize if blocksize > 0 else int(sample_rate // 10)
         self._stream_timeout_ms: int = stream_timeout_ms
         self._queue: asyncio.Queue[Optional[AudioFrame]] = asyncio.Queue(maxsize=capacity)
-        self._closed: bool = False
+        # _ending signals that no new streams will be added,
+        # but we continue processing until all streams are exhausted.
+        self._ending: bool = False
         self._mixer_task: asyncio.Task = asyncio.create_task(self._mixer())
 
     def add_stream(self, stream: AsyncIterator[AudioFrame]) -> None:
@@ -87,25 +88,10 @@ def remove_stream(self, stream: AsyncIterator[AudioFrame]) -> None:
         self._streams.discard(stream)
         self._buffers.pop(stream, None)
 
-    async def __aiter__(self) -> "AudioMixer":
-        """
-        Return the async iterator interface for the mixer.
-
-        Returns:
-            AudioMixer: The mixer itself, as it implements the async iterator protocol.
-        """
+    def __aiter__(self) -> "AudioMixer":
         return self
 
     async def __anext__(self) -> AudioFrame:
-        """
-        Retrieve the next mixed AudioFrame from the output queue.
-
-        Returns:
-            AudioFrame: The next mixed audio frame.
-
-        Raises:
-            StopAsyncIteration: When the mixer is closed and no more frames are available.
-        """
         item = await self._queue.get()
         if item is None:
             raise StopAsyncIteration
@@ -117,7 +103,7 @@ async def aclose(self) -> None:
 
         This cancels the mixing task, and any unconsumed output in the queue may be dropped.
         """
-        self._closed = True
+        self._ending = True
         self._mixer_task.cancel()
         with contextlib.suppress(asyncio.CancelledError):
             await self._mixer_task
@@ -129,13 +115,18 @@ def end_input(self) -> None:
         This method marks the mixer as closed so that it flushes any remaining buffered output before ending.
         Note that existing streams will still be processed until exhausted.
         """
-        self._closed = True
+        self._ending = True
 
     async def _mixer(self) -> None:
-        while not self._closed:
+        while True:
+            # If we're in ending mode and there are no more streams, exit.
+            if self._ending and not self._streams:
+                break
+
             if not self._streams:
                 await asyncio.sleep(0.01)
                 continue
+
             tasks = [
                 self._get_contribution(
                     stream,
@@ -185,7 +176,7 @@ async def _get_contribution(
                     stream.__anext__(), timeout=self._stream_timeout_ms / 1000
                 )
             except asyncio.TimeoutError:
-                logger.warning(f"AudioMixer: stream {stream} timeout, ignoring`")
+                logger.warning(f"AudioMixer: stream {stream} timeout, ignoring")
                 break
             except StopAsyncIteration:
                 exhausted = True

livekit-rtc/setup.py

@@ -58,7 +58,7 @@ def finalize_options(self):
     license="Apache-2.0",
     packages=setuptools.find_namespace_packages(include=["livekit.*"]),
     python_requires=">=3.9.0",
-    install_requires=["protobuf>=4.25.0", "types-protobuf>=3", "aiofiles>=24"],
+    install_requires=["protobuf>=4.25.0", "types-protobuf>=3", "aiofiles>=24", "numpy>=1.26"],
     package_data={
         "livekit.rtc": ["_proto/*.py", "py.typed", "*.pyi", "**/*.pyi"],
         "livekit.rtc.resources": ["*.so", "*.dylib", "*.dll", "LICENSE.md", "*.h"],

livekit-rtc/tests/test_mixer.py

@@ -0,0 +1,85 @@
+from typing import AsyncIterator
+import numpy as np
+import pytest
+import matplotlib.pyplot as plt
+
+from livekit.rtc import AudioMixer
+from livekit.rtc.audio_frame import AudioFrame
+
+SAMPLE_RATE = 48000
+# Use 100ms blocks (i.e. 1600 samples per frame)
+BLOCKSIZE = SAMPLE_RATE // 10
+
+
+async def sine_wave_generator(freq: float, duration: float) -> AsyncIterator[AudioFrame]:
+    total_frames = int((duration * SAMPLE_RATE) // BLOCKSIZE)
+    t_frame = np.arange(BLOCKSIZE) / SAMPLE_RATE
+    for i in range(total_frames):
+        # Shift the time for each frame so that the sine wave is continuous
+        t = t_frame + i * BLOCKSIZE / SAMPLE_RATE
+        # Create a sine wave with amplitude 0.3 (to avoid clipping when summing)
+        signal = 0.3 * np.sin(2 * np.pi * freq * t)
+        # Convert from float [-0.5, 0.5] to int16 values
+        signal_int16 = np.int16(signal * 32767)
+        frame = AudioFrame(
+            signal_int16.tobytes(),
+            SAMPLE_RATE,
+            1,
+            BLOCKSIZE,
+        )
+        yield frame
+
+
+@pytest.mark.asyncio
+async def test_mixer_two_sine_waves():
+    """
+    Test that mixing two sine waves (440Hz and 880Hz) produces an output
+    containing both frequency components.
+    """
+    duration = 1.0
+    mixer = AudioMixer(
+        sample_rate=SAMPLE_RATE,
+        num_channels=1,
+        blocksize=BLOCKSIZE,
+        stream_timeout_ms=100,
+        capacity=100,
+    )
+    stream1 = sine_wave_generator(440, duration)
+    stream2 = sine_wave_generator(880, duration)
+    mixer.add_stream(stream1)
+    mixer.add_stream(stream2)
+    mixer.end_input()
+
+    mixed_signals = []
+    async for frame in mixer:
+        data = np.frombuffer(frame.data.tobytes(), dtype=np.int16)
+        mixed_signals.append(data)
+
+    await mixer.aclose()
+
+    if not mixed_signals:
+        pytest.fail("No frames were produced by the mixer.")
+
+    mixed_signal = np.concatenate(mixed_signals)
+
+    plt.figure(figsize=(10, 4))
+    plt.plot(mixed_signal[:1000])  # plot 1000
+    plt.title("Mixed Signal")
+    plt.xlabel("Sample")
+    plt.ylabel("Amplitude")
+    plt.show()
+
+    # Use FFT to analyze frequency components.
+    fft = np.fft.rfft(mixed_signal)
+    freqs = np.fft.rfftfreq(len(mixed_signal), 1 / SAMPLE_RATE)
+    magnitude = np.abs(fft)
+
+    # Identify peak frequencies. We'll pick the 5 highest peaks.
+    peak_indices = np.argsort(magnitude)[-5:]
+    peak_freqs = freqs[peak_indices]
+
+    print("Peak frequencies:", peak_freqs)
+
+    # Assert that the peaks include 440Hz and 880Hz (with a tolerance of ±5 Hz)
+    assert any(np.isclose(peak_freqs, 440, atol=5)), f"Expected 440 Hz in peaks, got: {peak_freqs}"
+    assert any(np.isclose(peak_freqs, 880, atol=5)), f"Expected 880 Hz in peaks, got: {peak_freqs}"