120 lines
4.4 KiB
Python
120 lines
4.4 KiB
Python
"""vocal isolation: flag off, unavailable, success, timeout, cleanup (Demucs mocked)."""
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest import mock
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
|
|
|
from pydub import AudioSegment
|
|
from pydub.generators import Sine
|
|
|
|
from core import vocal_isolation as vi
|
|
from core.registry import TEMP_DIR
|
|
|
|
|
|
def wav(path, ms=1000):
|
|
Sine(300).to_audio_segment(ms).export(str(path), format="wav")
|
|
return path
|
|
|
|
|
|
class FakeProc:
|
|
"""Stands in for demucs: writes <out>/htdemucs/<stem>/vocals.wav, then exits."""
|
|
def __init__(self, args, **kw):
|
|
self.never_ends = FakeProc.never_ends
|
|
out = Path(args[args.index("-o") + 1])
|
|
if not self.never_ends and FakeProc.write:
|
|
d = out / "htdemucs" / Path(args[-1]).stem
|
|
d.mkdir(parents=True)
|
|
wav(d / "vocals.wav")
|
|
self.returncode = None if self.never_ends else FakeProc.rc
|
|
never_ends = False
|
|
write = True
|
|
rc = 0
|
|
def poll(self):
|
|
return self.returncode
|
|
def kill(self):
|
|
self.returncode = -9
|
|
def wait(self):
|
|
return self.returncode
|
|
|
|
|
|
def env(**kw):
|
|
base = {"VOCAL_ISOLATION": "1", "VOCAL_ISOLATION_CMD": "/bin/true"}
|
|
base.update(kw)
|
|
return mock.patch.dict(os.environ, base)
|
|
|
|
|
|
class VocalIsolationTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.src = wav(Path(tempfile.mkdtemp()) / "in.wav")
|
|
FakeProc.never_ends, FakeProc.write, FakeProc.rc = False, True, 0
|
|
self.w = mock.patch.object(vi, "_weights_present", return_value=True)
|
|
self.w.start()
|
|
self.addCleanup(self.w.stop)
|
|
|
|
def test_flag_off(self):
|
|
with mock.patch.dict(os.environ, {"VOCAL_ISOLATION": ""}):
|
|
self.assertFalse(vi.available())
|
|
with self.assertRaises(vi.VocalIsolationError) as c:
|
|
vi.isolate_vocals(self.src)
|
|
self.assertEqual(c.exception.status, 503)
|
|
|
|
def test_unavailable_without_weights_or_binary(self):
|
|
with env(), mock.patch.object(vi, "_weights_present", return_value=False):
|
|
self.assertFalse(vi.available())
|
|
with self.assertRaises(vi.VocalIsolationError) as c:
|
|
vi.isolate_vocals(self.src)
|
|
self.assertEqual(c.exception.status, 503)
|
|
with env(VOCAL_ISOLATION_CMD=""), mock.patch("shutil.which", return_value=None):
|
|
self.assertFalse(vi.available())
|
|
|
|
def test_success_and_cleanup(self):
|
|
with env(), mock.patch("subprocess.Popen", FakeProc):
|
|
self.assertTrue(vi.available())
|
|
out = vi.isolate_vocals(self.src)
|
|
self.assertTrue(out.exists())
|
|
self.assertEqual(AudioSegment.from_file(str(out)).frame_rate, 24000)
|
|
self.assertTrue(self.src.exists()) # original untouched
|
|
self.assertEqual(list(TEMP_DIR.glob("vocal_*")), [])
|
|
self.assertEqual(list(TEMP_DIR.glob("*_vocals.*")), [])
|
|
|
|
def test_timeout_kills_and_cleans(self):
|
|
FakeProc.never_ends = True
|
|
with env(VOCAL_ISOLATION_TIMEOUT="0.3"), mock.patch("subprocess.Popen", FakeProc):
|
|
with self.assertRaises(vi.VocalIsolationError) as c:
|
|
vi.isolate_vocals(self.src)
|
|
self.assertEqual(c.exception.status, 504)
|
|
self.assertEqual(list(TEMP_DIR.glob("vocal_*")), [])
|
|
|
|
def test_cancel(self):
|
|
import threading
|
|
FakeProc.never_ends = True
|
|
ev = threading.Event(); ev.set()
|
|
with env(), mock.patch("subprocess.Popen", FakeProc):
|
|
with self.assertRaises(vi.VocalIsolationError) as c:
|
|
vi.isolate_vocals(self.src, ev)
|
|
self.assertEqual(c.exception.status, 499)
|
|
|
|
def test_too_long_and_failure_and_busy(self):
|
|
with env(VOCAL_ISOLATION_MAX_SECONDS="0.5"):
|
|
with self.assertRaises(vi.VocalIsolationError) as c:
|
|
vi.isolate_vocals(self.src)
|
|
self.assertEqual(c.exception.status, 413)
|
|
FakeProc.write, FakeProc.rc = False, 1
|
|
with env(), mock.patch("subprocess.Popen", FakeProc):
|
|
with self.assertRaises(vi.VocalIsolationError) as c:
|
|
vi.isolate_vocals(self.src)
|
|
self.assertEqual(c.exception.status, 500)
|
|
self.assertEqual(list(TEMP_DIR.glob("vocal_*")), [])
|
|
with env(), vi._busy:
|
|
with self.assertRaises(vi.VocalIsolationError) as c:
|
|
vi.isolate_vocals(self.src)
|
|
self.assertEqual(c.exception.status, 409)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|