From 826af263b3a9f3da78a81dd821081a0dab39bc82 Mon Sep 17 00:00:00 2001
From: Joel Mathew Thomas <90510078+joelmathewthomas@users.noreply.github.com>
Date: Thu, 26 Dec 2024 20:46:37 +0530
Subject: [PATCH] added tests for demucs-wrapper.py, update deps

---
 .gitignore               |  3 ---
 requirements.txt         |  1 +
 tests/test_input.py      |  2 +-
 tests/test_separation.py | 53 ++++++++++++++++++++++++++++++++++++++++
 4 files changed, 55 insertions(+), 4 deletions(-)
 create mode 100644 tests/test_separation.py

diff --git a/.gitignore b/.gitignore
index ca64584..06b7abc 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,6 +14,3 @@ venv/
 
 # Ignore VSCode config
 .vscode/
-
-# Ignore output folders
-separated_output/
diff --git a/requirements.txt b/requirements.txt
index 5388cbf..27debed 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -84,6 +84,7 @@ requests==2.32.3
 retrying==1.3.4
 rich==13.9.4
 scikit-learn==1.6.0
+scipy==1.14.1
 setuptools==75.6.0
 six==1.17.0
 soundfile==0.12.1
diff --git a/tests/test_input.py b/tests/test_input.py
index afb4ff0..6c1c9c6 100644
--- a/tests/test_input.py
+++ b/tests/test_input.py
@@ -1,5 +1,5 @@
 import pytest
-from  src.input.file_reader import read_audio
+from src.input.file_reader import read_audio
 from src.input.format_checker import is_supported_format
 
 def test_read_audio():
diff --git a/tests/test_separation.py b/tests/test_separation.py
new file mode 100644
index 0000000..a297292
--- /dev/null
+++ b/tests/test_separation.py
@@ -0,0 +1,53 @@
+import os
+import pytest
+import tempfile
+import soundfile as sf
+from pathlib import Path
+from src.input.file_reader import read_audio
+from src.preprocessing.trim import trim_audio
+from src.preprocessing.resample import resample
+from src.separation.demucs_wrapper import separate_audio_with_demucs
+
+
+def test_demucs_separation_with_preprocessing():
+    """
+    Test to ensure Demucs separation works with preprocessing and creates expected outputs.
+    """
+
+    input_file = "./samples/am_contra_heart_peripheral.wav"
+    file_name = Path(input_file).stem
+    output_path = "/tmp/demucs-test"
+    waveform, samplerate = read_audio(input_file)
+    waveform = trim_audio(waveform, samplerate)
+
+    # Resample to 41000Hz
+    if samplerate != 41000:
+        print("Resampling audio to 41Khz")
+        waveform, samplerate = resample(waveform, samplerate, 41000)
+
+    # Create a temporary file to save the processed audio
+    temp_audio_path = tempfile.mktemp(suffix=".wav")
+
+    # Save the processed audio to the temporary file
+    sf.write(temp_audio_path, waveform, samplerate)
+
+    # Rename the file to orignal name
+    dir_path = os.path.dirname(temp_audio_path)
+    new_audio_path = os.path.join(dir_path, f"{file_name}.wav")
+    os.rename(temp_audio_path, new_audio_path)
+
+    separate_audio_with_demucs(new_audio_path, samplerate, output_path)
+
+    # Verify the htdemucs folder exists
+    demucs_dir = Path(output_path) / 'htdemucs'
+    assert demucs_dir.exists(), "htdemucs directory not found in output path."
+
+    # Verify the folder named after the file name (without extension) exists
+    file_folder = demucs_dir / file_name
+    assert file_folder.exists(), f"Folder {file_name} not found inside htdemucs directory."
+
+    # Verify the expected files exist inside the folder
+    expected_files = ['bass.wav', 'drums.wav', 'other.wav', 'vocals.wav']
+    for expected_file in expected_files:
+        file_path = file_folder / expected_file
+        assert file_path.exists(), f"Expected file {expected_file} not found in {file_name} folder."