Add fft_analysis/tests/test_core.py

fft_analysis/tests/test_core.py

@@ -0,0 +1,57 @@
+import pytest
+import numpy as np
+from math import isclose
+
+from src.fft_analysis.core import perform_fft
+
+
+def generate_sine_wave(freq: float, sample_rate: int, duration: float) -> np.ndarray:
+    t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
+    signal = np.sin(2 * np.pi * freq * t)
+    return signal
+
+
+def test_perform_fft_single_tone():
+    sample_rate = 8000  # Sampling rate in Hz
+    freq = 440.0  # A4 tone
+    duration = 1.0
+    audio_data = generate_sine_wave(freq, sample_rate, duration)
+
+    spectrum = perform_fft(audio_data)
+
+    # Expect spectrum length equals len(audio_data)//2 or similar
+    assert isinstance(spectrum, np.ndarray)
+    assert spectrum.ndim == 1
+    assert spectrum.size > 0
+
+    # Find peak frequency index
+    freq_bins = np.fft.rfftfreq(len(audio_data), d=1.0 / sample_rate)
+    peak_index = np.argmax(spectrum)
+    peak_freq = freq_bins[peak_index]
+
+    # Check if peak frequency corresponds roughly to the sine frequency
+    assert isclose(peak_freq, freq, rel_tol=0.02)
+
+
+def test_perform_fft_multiple_tones():
+    sample_rate = 8000
+    duration = 1.0
+    f1, f2 = 300.0, 900.0
+    signal = generate_sine_wave(f1, sample_rate, duration) + generate_sine_wave(f2, sample_rate, duration)
+
+    spectrum = perform_fft(signal)
+    freq_bins = np.fft.rfftfreq(len(signal), d=1.0 / sample_rate)
+    peaks = np.argsort(spectrum)[-2:]
+    peak_freqs = sorted(freq_bins[i] for i in peaks)
+
+    assert isclose(peak_freqs[0], f1, rel_tol=0.03) or isclose(peak_freqs[0], f2, rel_tol=0.03)
+    assert isclose(peak_freqs[1], f1, rel_tol=0.03) or isclose(peak_freqs[1], f2, rel_tol=0.03)
+
+
+def test_perform_fft_invalid_input():
+    # Expect error or empty spectrum for invalid input type
+    with pytest.raises((TypeError, ValueError)):
+        perform_fft('invalid_input')
+
+    with pytest.raises((TypeError, ValueError)):
+        perform_fft(np.array([]))