test_local_music20_acr.py 1.12 KB
import tempfile
import unittest
from pathlib import Path

import test_bootstrap

from scripts.local_music20_acr import collect_pairs, first_file


class LocalMusic20AcrTests(unittest.TestCase):
    def test_first_file_returns_none_for_missing_dir(self):
        with tempfile.TemporaryDirectory() as tmp:
            self.assertIsNone(first_file(Path(tmp) / 'missing'))

    def test_collect_pairs_uses_reference_and_query_types(self):
        with tempfile.TemporaryDirectory() as tmp:
            root = Path(tmp)
            song = root / '1'
            (song / 'type_11').mkdir(parents=True)
            (song / 'type_7').mkdir(parents=True)
            ((song / 'type_11') / 'ref.wav').write_bytes(b'r')
            ((song / 'type_7') / 'query.mp3').write_bytes(b'q')

            pairs = collect_pairs(root, song_limit=20, query_type=7)

            self.assertEqual(len(pairs), 1)
            self.assertEqual(pairs[0]['song_id'], '1')
            self.assertTrue(pairs[0]['reference_path'].endswith('ref.wav'))
            self.assertTrue(pairs[0]['query_path'].endswith('query.mp3'))


if __name__ == '__main__':
    unittest.main()