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,37 +10,39 @@ __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 for instrument in instruments:
for instrument in instruments: 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( TEST_AUDIO_DESCRIPTOR,
TEST_AUDIO_DESCRIPTOR, directory)
directory) for instrument in instruments:
for instrument in instruments: assert exists(join(directory, '{}.wav'.format(instrument)))
assert exists(join(directory, '{}.wav'.format(instrument))) # TODO: Consider testing generated file as well.
# TODO: Consider testing generated file as well.