diff --git a/fft_analysis/tests/test_core.py b/fft_analysis/tests/test_core.py new file mode 100644 index 0000000..a859fda --- /dev/null +++ b/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([]))