MusicUtils / tests /test_processor.py
sathvik77's picture
Add admin dashboard with download history and persistent storage
3b1b3cb
Raw
History Blame Contribute Delete
3.59 kB
import asyncio
from pathlib import Path
import pytest
from models import JobState
pytestmark = pytest.mark.asyncio
async def test_run_demucs_updates_state_on_success(tmp_path, monkeypatch):
"""Worker updates job from pending -> processing -> done (Demucs mocked)."""
from processor import _jobs, _run_demucs, create_job
job_id = "test-success-job"
input_path = tmp_path / "input.mp3"
output_dir = tmp_path / "outputs" / job_id
output_dir.mkdir(parents=True, exist_ok=True)
input_path.write_bytes(b"\xff\xfb" + b"\x00" * 100)
(output_dir / "original.mp3").write_bytes(b"\xff\xfb" + b"\x00" * 100)
class FakeProc:
def __init__(self, args):
self.args = args
self.returncode = 0
async def communicate(self):
if self.args[0] == "ffmpeg":
Path(self.args[-1]).parent.mkdir(parents=True, exist_ok=True)
Path(self.args[-1]).write_bytes(b"\xff\xfb" + b"\x00" * 50)
return b"", b""
# Simulate Demucs writing stems into the expected folder layout.
out_idx = self.args.index("-o")
demucs_work_dir = Path(self.args[out_idx + 1])
input_stem = Path(self.args[-1]).stem
base = demucs_work_dir / "htdemucs" / input_stem
base.mkdir(parents=True, exist_ok=True)
for stem_name in ("vocals", "drums", "bass", "other"):
(base / f"{stem_name}.mp3").write_bytes(b"\xff\xfb" + b"\x00" * 40)
return b"", b""
async def fake_create_subprocess(*args, **kwargs):
return FakeProc(args)
monkeypatch.setattr(asyncio, "create_subprocess_exec", fake_create_subprocess)
monkeypatch.setattr(
"processor.analyze_audio_file",
lambda _: {
"bpm": 128.4,
"bpm_confidence": 0.91,
"musical_key": "A",
"key_scale": "minor",
"key_confidence": 0.88,
},
)
create_job(job_id)
await _run_demucs(job_id, str(input_path), str(output_dir))
assert _jobs[job_id].state == JobState.done
assert _jobs[job_id].progress == 100
assert _jobs[job_id].stems is not None
assert set(_jobs[job_id].stems.keys()) == {"instrumental", "vocals", "drums", "bass", "other"}
for stem in _jobs[job_id].stems.values():
assert (output_dir / stem).exists()
assert _jobs[job_id].bpm == pytest.approx(128.4)
assert _jobs[job_id].bpm_confidence == pytest.approx(0.91)
assert _jobs[job_id].musical_key == "A"
assert _jobs[job_id].key_scale == "minor"
assert _jobs[job_id].key_confidence == pytest.approx(0.88)
async def test_run_demucs_marks_error_on_failure(tmp_path, monkeypatch):
"""Worker marks job as error when Demucs subprocess fails."""
from processor import _jobs, _run_demucs, create_job
job_id = "test-fail-job"
input_path = tmp_path / "input.mp3"
output_dir = tmp_path / "outputs" / job_id
input_path.write_bytes(b"\xff\xfb" + b"\x00" * 100)
class FakeFailProc:
returncode = 1
async def communicate(self):
return b"", b"Demucs error: model not found"
async def fake_create_subprocess(*args, **kwargs):
return FakeFailProc()
monkeypatch.setattr(asyncio, "create_subprocess_exec", fake_create_subprocess)
create_job(job_id)
with pytest.raises(RuntimeError, match="Demucs failed"):
await _run_demucs(job_id, str(input_path), str(output_dir))
assert _jobs[job_id].state == JobState.error
assert _jobs[job_id].error is not None