fix: test parameters

This commit is contained in:
Félix Voituret
2019-11-08 20:03:16 -05:00
parent 05571be661
commit e9f1ba9450

View File

@@ -10,22 +10,24 @@ __license__ = 'MIT License'
from os.path import exists, join from os.path import exists, join
from tempfile import TemporaryDirectory from tempfile import TemporaryDirectory
import pytest
from spleeter.audio.adapter import get_default_audio_adapter from spleeter.audio.adapter import get_default_audio_adapter
from spleeter.separator import Separator from spleeter.separator import Separator
TEST_AUDIO_DESCRIPTOR = 'audio_example.mp3' TEST_AUDIO_DESCRIPTOR = 'audio_example.mp3'
TEST_CONFIGURATIONS = { TEST_CONFIGURATIONS = [
'spleeter:2stems': ('vocals', 'accompaniment'), ('spleeter:2stems', ('vocals', 'accompaniment')),
'spleeter:4stems': ('vocals', 'drums', 'bass', 'other'), ('spleeter:4stems', ('vocals', 'drums', 'bass', 'other')),
'spleeter:5stems': ('vocals', 'drums', 'bass', 'piano', 'other') ('spleeter:5stems', ('vocals', 'drums', 'bass', 'piano', 'other'))
} ]
@pytest.mark.parametrize('configuration, instruments', TEST_CONFIGURATIONS)
def test_separate(): def test_separate():
""" Test separation from raw data. """ """ Test separation from raw data. """
adapter = get_default_audio_adapter() adapter = get_default_audio_adapter()
waveform, _ = adapter.load(TEST_AUDIO_DESCRIPTOR) waveform, _ = adapter.load(TEST_AUDIO_DESCRIPTOR)
for configuration, instruments in TEST_CONFIGURATIONS.items():
separator = Separator(configuration) separator = Separator(configuration)
prediction = separator.separate(waveform) prediction = separator.separate(waveform)
assert len(prediction) == 2 assert len(prediction) == 2
@@ -33,9 +35,9 @@ def test_separate():
assert instrument in prediction assert instrument in prediction
def test_separate_to_file(): @pytest.mark.parametrize('configuration, instruments', TEST_CONFIGURATIONS)
def test_separate_to_file(configuration, instruments):
""" Test file based separation. """ """ Test file based separation. """
for configuration, instruments in TEST_CONFIGURATIONS.items():
separator = Separator(configuration) separator = Separator(configuration)
with TemporaryDirectory() as directory: with TemporaryDirectory() as directory:
separator.separate_to_file( separator.separate_to_file(