annotate tests/test_io.py @ 144:fae263792c2e

Add tests for _normalize_use_channel
author Amine Sehili <amine.sehili@gmail.com>
date Fri, 18 Feb 2028 16:10:48 +0100
parents 5ff29f286374
children 85d26e96d259
rev   line source
amine@106 1 import os
amine@106 2 import sys
amine@106 3 import math
amine@107 4 from array import array
amine@133 5 from tempfile import NamedTemporaryFile, TemporaryDirectory
amine@110 6 import filecmp
amine@108 7 from unittest import TestCase
amine@108 8 from genty import genty, genty_dataset
amine@110 9 from auditok.io import (
amine@126 10 DATA_FORMAT,
amine@121 11 AudioIOError,
amine@110 12 AudioParameterError,
amine@126 13 BufferAudioSource,
amine@110 14 check_audio_data,
amine@143 15 _guess_audio_format,
amine@144 16 _normalize_use_channel,
amine@128 17 _get_audio_parameters,
amine@116 18 _array_to_bytes,
amine@118 19 _mix_audio_channels,
amine@119 20 _extract_selected_channel,
amine@126 21 _load_raw,
amine@129 22 _load_wave,
amine@131 23 _load_with_pydub,
amine@120 24 from_file,
amine@111 25 _save_raw,
amine@110 26 _save_wave,
amine@141 27 _save_with_pydub,
amine@135 28 to_file,
amine@110 29 )
amine@106 30
amine@106 31
amine@106 32 if sys.version_info >= (3, 0):
amine@106 33 PYTHON_3 = True
amine@124 34 from unittest.mock import patch, Mock
amine@106 35 else:
amine@106 36 PYTHON_3 = False
amine@124 37 from mock import patch, Mock
amine@120 38
amine@120 39 AUDIO_PARAMS_SHORT = {"sr": 16000, "sw": 2, "ch": 1}
amine@106 40
amine@106 41
amine@106 42 def _sample_generator(*data_buffers):
amine@106 43 """
amine@106 44 Takes a list of many mono audio data buffers and makes a sample generator
amine@106 45 of interleaved audio samples, one sample from each channel. The resulting
amine@106 46 generator can be used to build a multichannel audio buffer.
amine@106 47 >>> gen = _sample_generator("abcd", "ABCD")
amine@106 48 >>> list(gen)
amine@106 49 ["a", "A", "b", "B", "c", "C", "d", "D"]
amine@106 50 """
amine@106 51 frame_gen = zip(*data_buffers)
amine@106 52 return (sample for frame in frame_gen for sample in frame)
amine@106 53
amine@106 54
amine@107 55 def _generate_pure_tone(
amine@107 56 frequency, duration_sec=1, sampling_rate=16000, sample_width=2, volume=1e4
amine@107 57 ):
amine@107 58 """
amine@107 59 Generates a pure tone with the given frequency.
amine@107 60 """
amine@107 61 assert frequency <= sampling_rate / 2
amine@107 62 max_value = (2 ** (sample_width * 8) // 2) - 1
amine@107 63 if volume > max_value:
amine@107 64 volume = max_value
amine@107 65 fmt = DATA_FORMAT[sample_width]
amine@107 66 total_samples = int(sampling_rate * duration_sec)
amine@107 67 step = frequency / sampling_rate
amine@107 68 two_pi_step = 2 * math.pi * step
amine@107 69 data = array(
amine@107 70 fmt,
amine@107 71 (
amine@107 72 int(math.sin(two_pi_step * i) * volume)
amine@107 73 for i in range(total_samples)
amine@107 74 ),
amine@107 75 )
amine@107 76 return data
amine@107 77
amine@107 78
amine@107 79 PURE_TONE_DICT = {
amine@107 80 freq: _generate_pure_tone(freq, 1, 16000, 2) for freq in (400, 800, 1600)
amine@107 81 }
amine@107 82 PURE_TONE_DICT.update(
amine@107 83 {
amine@107 84 freq: _generate_pure_tone(freq, 0.1, 16000, 2)
amine@107 85 for freq in (600, 1150, 2400, 7220)
amine@107 86 }
amine@107 87 )
amine@108 88
amine@108 89
amine@108 90 @genty
amine@108 91 class TestIO(TestCase):
amine@108 92 @genty_dataset(
amine@108 93 valid_mono=(b"\0" * 113, 1, 1),
amine@108 94 valid_stereo=(b"\0" * 160, 1, 2),
amine@108 95 invalid_mono_sw_2=(b"\0" * 113, 2, 1, False),
amine@108 96 invalid_stereo_sw_1=(b"\0" * 113, 1, 2, False),
amine@108 97 invalid_stereo_sw_2=(b"\0" * 158, 2, 2, False),
amine@108 98 )
amine@108 99 def test_check_audio_data(self, data, sample_width, channels, valid=True):
amine@108 100
amine@108 101 if not valid:
amine@108 102 with self.assertRaises(AudioParameterError):
amine@108 103 check_audio_data(data, sample_width, channels)
amine@108 104 else:
amine@108 105 self.assertIsNone(check_audio_data(data, sample_width, channels))
amine@110 106
amine@110 107 @genty_dataset(
amine@143 108 extention_and_format_same=("wav", "filename.wav", "wav"),
amine@143 109 extention_and_format_different=("wav", "filename.mp3", "wav"),
amine@143 110 extention_no_format=(None, "filename.wav", "wav"),
amine@143 111 format_no_extension=("wav", "filename", "wav"),
amine@143 112 no_format_no_extension=(None, "filename", None),
amine@143 113 )
amine@143 114 def test_guess_audio_format(self, fmt, filename, expected):
amine@143 115 result = _guess_audio_format(fmt, filename)
amine@143 116 self.assertEqual(result, expected)
amine@143 117
amine@143 118 @genty_dataset(
amine@144 119 none=(None, 0),
amine@144 120 positive_int=(1, 1),
amine@144 121 negative_int=(-1, -1),
amine@144 122 left=("left", 0),
amine@144 123 right=("right", 1),
amine@144 124 mix=("mix", "mix"),
amine@144 125 )
amine@144 126 def test_normalize_use_channel(self, use_channel, expected):
amine@144 127 result = _normalize_use_channel(use_channel)
amine@144 128 self.assertEqual(result, expected)
amine@144 129
amine@144 130 @genty_dataset(
amine@118 131 mono_1byte=([400], 1),
amine@118 132 stereo_1byte=([400, 600], 1),
amine@118 133 three_channel_1byte=([400, 600, 2400], 1),
amine@118 134 mono_2byte=([400], 2),
amine@118 135 stereo_2byte=([400, 600], 2),
amine@118 136 three_channel_2byte=([400, 600, 1150], 2),
amine@118 137 mono_4byte=([400], 4),
amine@118 138 stereo_4byte=([400, 600], 4),
amine@118 139 four_channel_2byte=([400, 600, 1150, 7220], 4),
amine@118 140 )
amine@118 141 def test_mix_audio_channels(self, frequencies, sample_width):
amine@118 142 sampling_rate = 16000
amine@118 143 sample_width = 2
amine@118 144 channels = len(frequencies)
amine@118 145 mono_channels = [
amine@118 146 _generate_pure_tone(
amine@118 147 freq,
amine@118 148 duration_sec=0.1,
amine@118 149 sampling_rate=sampling_rate,
amine@118 150 sample_width=sample_width,
amine@118 151 )
amine@118 152 for freq in frequencies
amine@118 153 ]
amine@118 154 fmt = DATA_FORMAT[sample_width]
amine@118 155 expected = _array_to_bytes(
amine@118 156 array(
amine@118 157 fmt,
amine@118 158 (sum(samples) // channels for samples in zip(*mono_channels)),
amine@118 159 )
amine@118 160 )
amine@118 161 data = _array_to_bytes(array(fmt, _sample_generator(*mono_channels)))
amine@118 162 mixed = _mix_audio_channels(data, channels, sample_width)
amine@118 163 self.assertEqual(mixed, expected)
amine@118 164
amine@118 165 @genty_dataset(
amine@119 166 mono_1byte=([400], 1, 0),
amine@119 167 stereo_1byte_2st_channel=([400, 600], 1, 1),
amine@119 168 mono_2byte=([400], 2, 0),
amine@119 169 stereo_2byte_1st_channel=([400, 600], 2, 0),
amine@119 170 stereo_2byte_2nd_channel=([400, 600], 2, 1),
amine@119 171 three_channel_2byte_last_negative_idx=([400, 600, 1150], 2, -1),
amine@119 172 three_channel_2byte_2nd_negative_idx=([400, 600, 1150], 2, -2),
amine@119 173 three_channel_2byte_1st_negative_idx=([400, 600, 1150], 2, -3),
amine@119 174 three_channel_4byte_1st=([400, 600, 1150], 4, 0),
amine@119 175 three_channel_4byte_last_negative_idx=([400, 600, 1150], 4, -1),
amine@119 176 )
amine@119 177 def test_extract_selected_channel(
amine@119 178 self, frequencies, sample_width, use_channel
amine@119 179 ):
amine@119 180
amine@119 181 mono_channels = [
amine@119 182 _generate_pure_tone(
amine@119 183 freq,
amine@119 184 duration_sec=0.1,
amine@119 185 sampling_rate=16000,
amine@119 186 sample_width=sample_width,
amine@119 187 )
amine@119 188 for freq in frequencies
amine@119 189 ]
amine@119 190 channels = len(frequencies)
amine@119 191 fmt = DATA_FORMAT[sample_width]
amine@119 192 expected = _array_to_bytes(mono_channels[use_channel])
amine@119 193 data = _array_to_bytes(array(fmt, _sample_generator(*mono_channels)))
amine@119 194 selected_channel = _extract_selected_channel(
amine@119 195 data, channels, sample_width, use_channel
amine@119 196 )
amine@119 197 self.assertEqual(selected_channel, expected)
amine@119 198
amine@119 199 @genty_dataset(
amine@120 200 raw_with_audio_format=(
amine@120 201 "audio",
amine@120 202 "raw",
amine@120 203 "_load_raw",
amine@120 204 AUDIO_PARAMS_SHORT,
amine@120 205 ),
amine@120 206 raw_with_extension=(
amine@120 207 "audio.raw",
amine@120 208 None,
amine@120 209 "_load_raw",
amine@120 210 AUDIO_PARAMS_SHORT,
amine@120 211 ),
amine@120 212 wave_with_audio_format=("audio", "wave", "_load_wave"),
amine@120 213 wav_with_audio_format=("audio", "wave", "_load_wave"),
amine@120 214 wav_with_extension=("audio.wav", None, "_load_wave"),
amine@120 215 format_and_extension_both_given=("audio.dat", "wav", "_load_wave"),
amine@120 216 format_and_extension_both_given_b=("audio.raw", "wave", "_load_wave"),
amine@120 217 no_format_nor_extension=("audio", None, "_load_with_pydub"),
amine@120 218 other_formats_ogg=("audio.ogg", None, "_load_with_pydub"),
amine@120 219 other_formats_webm=("audio", "webm", "_load_with_pydub"),
amine@120 220 )
amine@120 221 def test_from_file(
amine@120 222 self, filename, audio_format, funtion_name, kwargs=None
amine@120 223 ):
amine@120 224 funtion_name = "auditok.io." + funtion_name
amine@120 225 if kwargs is None:
amine@120 226 kwargs = {}
amine@120 227 with patch(funtion_name) as patch_function:
amine@120 228 from_file(filename, audio_format, **kwargs)
amine@120 229 self.assertTrue(patch_function.called)
amine@120 230
amine@137 231 @genty_dataset(
amine@137 232 missing_sampling_rate=("sr",),
amine@137 233 missing_sample_width=("sw",),
amine@137 234 missing_channels=("ch",),
amine@137 235 )
amine@137 236 def test_from_file_missing_audio_param(self, missing_param):
amine@137 237 with self.assertRaises(AudioParameterError):
amine@137 238 params = AUDIO_PARAMS_SHORT.copy()
amine@137 239 del params[missing_param]
amine@137 240 from_file("audio", audio_format="raw", **params)
amine@137 241
amine@121 242 def test_from_file_no_pydub(self):
amine@121 243 with patch("auditok.io._WITH_PYDUB", False):
amine@121 244 with self.assertRaises(AudioIOError):
amine@121 245 from_file("audio", "mp3")
amine@121 246
amine@111 247 @genty_dataset(
amine@122 248 raw_first_channel=("raw", 0, 400),
amine@122 249 raw_second_channel=("raw", 1, 800),
amine@122 250 raw_third_channel=("raw", 2, 1600),
amine@122 251 raw_left_channel=("raw", "left", 400),
amine@122 252 raw_right_channel=("raw", "right", 800),
amine@122 253 wav_first_channel=("wav", 0, 400),
amine@122 254 wav_second_channel=("wav", 1, 800),
amine@122 255 wav_third_channel=("wav", 2, 1600),
amine@122 256 wav_left_channel=("wav", "left", 400),
amine@122 257 wav_right_channel=("wav", "right", 800),
amine@122 258 )
amine@122 259 def test_from_file_multichannel_audio(
amine@122 260 self, audio_format, use_channel, frequency
amine@122 261 ):
amine@122 262 expected = PURE_TONE_DICT[frequency]
amine@122 263 filename = "tests/data/test_16KHZ_3channel_400-800-1600Hz.{}".format(
amine@122 264 audio_format
amine@122 265 )
amine@122 266 sample_width = 2
amine@122 267 audio_source = from_file(
amine@122 268 filename,
amine@122 269 sampling_rate=16000,
amine@122 270 sample_width=sample_width,
amine@122 271 channels=3,
amine@122 272 use_channel=use_channel,
amine@122 273 )
amine@122 274 fmt = DATA_FORMAT[sample_width]
amine@122 275 data = array(fmt, audio_source._buffer)
amine@122 276 self.assertEqual(data, expected)
amine@122 277
amine@122 278 @genty_dataset(
amine@123 279 raw_mono=("raw", "mono_400Hz", (400,)),
amine@123 280 raw_3channel=("raw", "3channel_400-800-1600Hz", (400, 800, 1600)),
amine@123 281 wav_mono=("wav", "mono_400Hz", (400,)),
amine@123 282 wav_3channel=("wav", "3channel_400-800-1600Hz", (400, 800, 1600)),
amine@123 283 )
amine@123 284 def test_from_file_multichannel_audio_mix(
amine@123 285 self, audio_format, filename_suffix, frequencies
amine@123 286 ):
amine@123 287 sampling_rate = 16000
amine@123 288 sample_width = 2
amine@123 289 channels = len(frequencies)
amine@123 290 mono_channels = [PURE_TONE_DICT[freq] for freq in frequencies]
amine@123 291 channels = len(frequencies)
amine@123 292 fmt = DATA_FORMAT[sample_width]
amine@123 293 expected = _array_to_bytes(
amine@123 294 array(
amine@123 295 fmt,
amine@123 296 (sum(samples) // channels for samples in zip(*mono_channels)),
amine@123 297 )
amine@123 298 )
amine@123 299 filename = "tests/data/test_16KHZ_{}.{}".format(
amine@123 300 filename_suffix, audio_format
amine@123 301 )
amine@123 302 audio_source = from_file(
amine@123 303 filename,
amine@123 304 use_channel="mix",
amine@123 305 sampling_rate=sampling_rate,
amine@123 306 sample_width=2,
amine@123 307 channels=channels,
amine@123 308 )
amine@123 309 mixed = audio_source._buffer
amine@123 310 self.assertEqual((mixed), expected)
amine@123 311
amine@124 312 @patch("auditok.io._WITH_PYDUB", True)
amine@124 313 @patch("auditok.io.BufferAudioSource")
amine@124 314 @genty_dataset(
amine@124 315 ogg_first_channel=("ogg", 0, "from_ogg"),
amine@124 316 ogg_second_channel=("ogg", 1, "from_ogg"),
amine@124 317 ogg_mix=("ogg", "mix", "from_ogg"),
amine@124 318 ogg_default=("ogg", None, "from_ogg"),
amine@124 319 mp3_left_channel=("mp3", "left", "from_mp3"),
amine@124 320 mp3_right_channel=("mp3", "right", "from_mp3"),
amine@124 321 flac_first_channel=("flac", 0, "from_file"),
amine@124 322 flac_second_channel=("flac", 1, "from_file"),
amine@124 323 flv_left_channel=("flv", "left", "from_flv"),
amine@124 324 webm_right_channel=("webm", "right", "from_file"),
amine@124 325 )
amine@124 326 def test_from_file_multichannel_audio_compressed(
amine@124 327 self, audio_format, use_channel, function, *mocks
amine@124 328 ):
amine@124 329 filename = "audio.{}".format(audio_format)
amine@124 330 segment_mock = Mock()
amine@124 331 segment_mock.sample_width = 2
amine@124 332 segment_mock.channels = 2
amine@124 333 segment_mock._data = b"abcd"
amine@124 334 with patch("auditok.io._extract_selected_channel") as ext_mock:
amine@124 335 with patch(
amine@124 336 "auditok.io.AudioSegment.{}".format(function)
amine@124 337 ) as open_func:
amine@124 338 open_func.return_value = segment_mock
amine@124 339 from_file(filename, use_channel=use_channel)
amine@124 340 self.assertTrue(open_func.called)
amine@124 341 self.assertTrue(ext_mock.called)
amine@124 342
amine@124 343 use_channel = {"left": 0, "right": 1, None: 0}.get(
amine@124 344 use_channel, use_channel
amine@124 345 )
amine@124 346 ext_mock.assert_called_with(
amine@124 347 segment_mock._data,
amine@124 348 segment_mock.channels,
amine@124 349 segment_mock.sample_width,
amine@124 350 use_channel,
amine@124 351 )
amine@124 352
amine@124 353 with patch("auditok.io._extract_selected_channel") as ext_mock:
amine@124 354 with patch(
amine@124 355 "auditok.io.AudioSegment.{}".format(function)
amine@124 356 ) as open_func:
amine@124 357 segment_mock.channels = 1
amine@124 358 open_func.return_value = segment_mock
amine@124 359 from_file(filename, use_channel=use_channel)
amine@124 360 self.assertTrue(open_func.called)
amine@124 361 self.assertFalse(ext_mock.called)
amine@124 362
amine@125 363 @patch("auditok.io._WITH_PYDUB", True)
amine@125 364 @patch("auditok.io.BufferAudioSource")
amine@125 365 @genty_dataset(
amine@125 366 ogg=("ogg", "from_ogg"),
amine@125 367 mp3=("mp3", "from_mp3"),
amine@125 368 flac=("flac", "from_file"),
amine@125 369 )
amine@125 370 def test_from_file_multichannel_audio_mix_compressed(
amine@125 371 self, audio_format, function, *mocks
amine@125 372 ):
amine@125 373 filename = "audio.{}".format(audio_format)
amine@125 374 segment_mock = Mock()
amine@125 375 segment_mock.sample_width = 2
amine@125 376 segment_mock.channels = 2
amine@125 377 segment_mock._data = b"abcd"
amine@125 378 with patch("auditok.io._mix_audio_channels") as mix_mock:
amine@125 379 with patch(
amine@125 380 "auditok.io.AudioSegment.{}".format(function)
amine@125 381 ) as open_func:
amine@125 382 open_func.return_value = segment_mock
amine@125 383 from_file(filename, use_channel="mix")
amine@125 384 self.assertTrue(open_func.called)
amine@125 385 mix_mock.assert_called_with(
amine@125 386 segment_mock._data,
amine@125 387 segment_mock.channels,
amine@125 388 segment_mock.sample_width,
amine@125 389 )
amine@125 390
amine@123 391 @genty_dataset(
amine@126 392 dafault_first_channel=(None, 400),
amine@126 393 first_channel=(0, 400),
amine@126 394 second_channel=(1, 800),
amine@126 395 third_channel=(2, 1600),
amine@126 396 negative_first_channel=(-3, 400),
amine@126 397 negative_second_channel=(-2, 800),
amine@126 398 negative_third_channel=(-1, 1600),
amine@126 399 )
amine@126 400 def test_load_raw(self, use_channel, frequency):
amine@126 401 filename = "tests/data/test_16KHZ_3channel_400-800-1600Hz.raw"
amine@126 402 if use_channel is not None:
amine@126 403 audio_source = _load_raw(
amine@126 404 filename,
amine@126 405 sampling_rate=16000,
amine@126 406 sample_width=2,
amine@126 407 channels=3,
amine@126 408 use_channel=use_channel,
amine@126 409 )
amine@126 410 else:
amine@126 411 audio_source = _load_raw(
amine@126 412 filename, sampling_rate=16000, sample_width=2, channels=3
amine@126 413 )
amine@126 414 self.assertIsInstance(audio_source, BufferAudioSource)
amine@126 415 self.assertEqual(audio_source.sampling_rate, 16000)
amine@126 416 self.assertEqual(audio_source.sample_width, 2)
amine@126 417 self.assertEqual(audio_source.channels, 1)
amine@126 418 # generate a pure sine wave tone of the given frequency
amine@126 419 expected = PURE_TONE_DICT[frequency]
amine@126 420 # compre with data read from file
amine@126 421 fmt = DATA_FORMAT[2]
amine@126 422 data = array(fmt, audio_source._buffer)
amine@126 423 self.assertEqual(data, expected)
amine@126 424
amine@126 425 @genty_dataset(
amine@127 426 mono=("mono_400Hz", (400,)),
amine@127 427 three_channel=("3channel_400-800-1600Hz", (400, 800, 1600)),
amine@127 428 )
amine@127 429 def test_load_raw_mix(self, filename_suffix, frequencies):
amine@127 430 sampling_rate = 16000
amine@127 431 sample_width = 2
amine@127 432 channels = len(frequencies)
amine@127 433 mono_channels = [PURE_TONE_DICT[freq] for freq in frequencies]
amine@127 434
amine@127 435 fmt = DATA_FORMAT[sample_width]
amine@127 436 expected = _array_to_bytes(
amine@127 437 array(
amine@127 438 fmt,
amine@127 439 (sum(samples) // channels for samples in zip(*mono_channels)),
amine@127 440 )
amine@127 441 )
amine@127 442 filename = "tests/data/test_16KHZ_{}.raw".format(filename_suffix)
amine@127 443 audio_source = _load_raw(
amine@127 444 filename,
amine@127 445 use_channel="mix",
amine@127 446 sampling_rate=sampling_rate,
amine@127 447 sample_width=2,
amine@127 448 channels=channels,
amine@127 449 )
amine@127 450 mixed = audio_source._buffer
amine@127 451 self.assertEqual(mixed, expected)
amine@127 452 self.assertIsInstance(audio_source, BufferAudioSource)
amine@127 453 self.assertEqual(audio_source.sampling_rate, sampling_rate)
amine@127 454 self.assertEqual(audio_source.sample_width, sample_width)
amine@127 455 self.assertEqual(audio_source.channels, 1)
amine@127 456
amine@127 457 @genty_dataset(
amine@128 458 missing_sampling_rate=("sr",),
amine@128 459 missing_sample_width=("sw",),
amine@128 460 missing_channels=("ch",),
amine@128 461 )
amine@128 462 def test_load_raw_missing_audio_param(self, missing_param):
amine@128 463 with self.assertRaises(AudioParameterError):
amine@128 464 params = AUDIO_PARAMS_SHORT.copy()
amine@128 465 del params[missing_param]
amine@128 466 srate, swidth, channels, _ = _get_audio_parameters(params)
amine@128 467 _load_raw("audio", srate, swidth, channels)
amine@128 468
amine@128 469 @genty_dataset(
amine@129 470 dafault_first_channel=(None, 400),
amine@129 471 first_channel=(0, 400),
amine@129 472 second_channel=(1, 800),
amine@129 473 third_channel=(2, 1600),
amine@129 474 negative_first_channel=(-3, 400),
amine@129 475 negative_second_channel=(-2, 800),
amine@129 476 negative_third_channel=(-1, 1600),
amine@129 477 )
amine@129 478 def test_load_wave(self, use_channel, frequency):
amine@129 479 filename = "tests/data/test_16KHZ_3channel_400-800-1600Hz.wav"
amine@129 480 if use_channel is not None:
amine@129 481 audio_source = _load_wave(filename, use_channel=use_channel)
amine@129 482 else:
amine@129 483 audio_source = _load_wave(filename)
amine@129 484 self.assertIsInstance(audio_source, BufferAudioSource)
amine@129 485 self.assertEqual(audio_source.sampling_rate, 16000)
amine@129 486 self.assertEqual(audio_source.sample_width, 2)
amine@129 487 self.assertEqual(audio_source.channels, 1)
amine@129 488 # generate a pure sine wave tone of the given frequency
amine@129 489 expected = PURE_TONE_DICT[frequency]
amine@129 490 # compre with data read from file
amine@129 491 fmt = DATA_FORMAT[2]
amine@129 492 data = array(fmt, audio_source._buffer)
amine@129 493 self.assertEqual(data, expected)
amine@129 494
amine@129 495 @genty_dataset(
amine@130 496 mono=("mono_400Hz", (400,)),
amine@130 497 three_channel=("3channel_400-800-1600Hz", (400, 800, 1600)),
amine@130 498 )
amine@130 499 def test_load_wave_mix(self, filename_suffix, frequencies):
amine@130 500 sampling_rate = 16000
amine@130 501 sample_width = 2
amine@130 502 channels = len(frequencies)
amine@130 503 mono_channels = [PURE_TONE_DICT[freq] for freq in frequencies]
amine@130 504 fmt = DATA_FORMAT[sample_width]
amine@130 505 expected = _array_to_bytes(
amine@130 506 array(
amine@130 507 fmt,
amine@130 508 (sum(samples) // channels for samples in zip(*mono_channels)),
amine@130 509 )
amine@130 510 )
amine@130 511 filename = "tests/data/test_16KHZ_{}.wav".format(filename_suffix)
amine@130 512 audio_source = _load_wave(filename, use_channel="mix")
amine@130 513 mixed = audio_source._buffer
amine@130 514 self.assertEqual(mixed, expected)
amine@130 515 self.assertIsInstance(audio_source, BufferAudioSource)
amine@130 516 self.assertEqual(audio_source.sampling_rate, sampling_rate)
amine@130 517 self.assertEqual(audio_source.sample_width, sample_width)
amine@130 518 self.assertEqual(audio_source.channels, 1)
amine@130 519
amine@131 520 @patch("auditok.io._WITH_PYDUB", True)
amine@131 521 @patch("auditok.io.BufferAudioSource")
amine@131 522 @genty_dataset(
amine@131 523 ogg_default_first_channel=("ogg", 2, None, "from_ogg"),
amine@131 524 ogg_first_channel=("ogg", 1, 0, "from_ogg"),
amine@131 525 ogg_second_channel=("ogg", 2, 1, "from_ogg"),
amine@131 526 ogg_mix_channels=("ogg", 3, "mix", "from_ogg"),
amine@131 527 mp3_left_channel=("mp3", 1, "left", "from_mp3"),
amine@131 528 mp3_right_channel=("mp3", 2, "right", "from_mp3"),
amine@131 529 mp3_mix_channels=("mp3", 3, "mix", "from_mp3"),
amine@131 530 flac_first_channel=("flac", 2, 0, "from_file"),
amine@131 531 flac_second_channel=("flac", 2, 1, "from_file"),
amine@131 532 flv_left_channel=("flv", 1, "left", "from_flv"),
amine@131 533 webm_right_channel=("webm", 2, "right", "from_file"),
amine@131 534 webm_mix_channels=("webm", 4, "mix", "from_file"),
amine@131 535 )
amine@131 536 def test_load_with_pydub(
amine@131 537 self, audio_format, channels, use_channel, function, *mocks
amine@131 538 ):
amine@131 539 filename = "audio.{}".format(audio_format)
amine@131 540 segment_mock = Mock()
amine@131 541 segment_mock.sample_width = 2
amine@131 542 segment_mock.channels = channels
amine@131 543 segment_mock._data = b"abcdefgh"
amine@131 544 with patch("auditok.io._extract_selected_channel") as ext_mock:
amine@131 545 with patch(
amine@131 546 "auditok.io.AudioSegment.{}".format(function)
amine@131 547 ) as open_func:
amine@131 548 open_func.return_value = segment_mock
amine@131 549 use_channel = {"left": 0, "right": 1, None: 0}.get(
amine@131 550 use_channel, use_channel
amine@131 551 )
amine@131 552 _load_with_pydub(filename, audio_format, use_channel)
amine@131 553 self.assertTrue(open_func.called)
amine@131 554 if channels > 1:
amine@131 555 self.assertTrue(ext_mock.called)
amine@131 556 ext_mock.assert_called_with(
amine@131 557 segment_mock._data,
amine@131 558 segment_mock.channels,
amine@131 559 segment_mock.sample_width,
amine@131 560 use_channel,
amine@131 561 )
amine@131 562 else:
amine@131 563 self.assertFalse(ext_mock.called)
amine@131 564
amine@130 565 @genty_dataset(
amine@132 566 mono=("mono_400Hz.raw", (400,)),
amine@132 567 three_channel=("3channel_400-800-1600Hz.raw", (400, 800, 1600)),
amine@132 568 )
amine@132 569 def test_save_raw(self, filename, frequencies):
amine@132 570 filename = "tests/data/test_16KHZ_{}".format(filename)
amine@132 571 sample_width = 2
amine@132 572 fmt = DATA_FORMAT[sample_width]
amine@132 573 mono_channels = [PURE_TONE_DICT[freq] for freq in frequencies]
amine@132 574 data = _array_to_bytes(array(fmt, _sample_generator(*mono_channels)))
amine@132 575 tmpfile = NamedTemporaryFile()
amine@136 576 _save_raw(data, tmpfile.name)
amine@132 577 self.assertTrue(filecmp.cmp(tmpfile.name, filename, shallow=False))
amine@132 578
amine@132 579 @genty_dataset(
amine@110 580 mono=("mono_400Hz.wav", (400,)),
amine@110 581 three_channel=("3channel_400-800-1600Hz.wav", (400, 800, 1600)),
amine@110 582 )
amine@110 583 def test_save_wave(self, filename, frequencies):
amine@110 584 filename = "tests/data/test_16KHZ_{}".format(filename)
amine@110 585 sampling_rate = 16000
amine@110 586 sample_width = 2
amine@110 587 channels = len(frequencies)
amine@110 588 fmt = DATA_FORMAT[sample_width]
amine@110 589 mono_channels = [PURE_TONE_DICT[freq] for freq in frequencies]
amine@110 590 data = _array_to_bytes(array(fmt, _sample_generator(*mono_channels)))
amine@110 591 tmpfile = NamedTemporaryFile()
amine@136 592 _save_wave(data, tmpfile.name, sampling_rate, sample_width, channels)
amine@110 593 self.assertTrue(filecmp.cmp(tmpfile.name, filename, shallow=False))
amine@132 594
amine@132 595 @genty_dataset(
amine@132 596 missing_sampling_rate=("sr",),
amine@132 597 missing_sample_width=("sw",),
amine@132 598 missing_channels=("ch",),
amine@132 599 )
amine@132 600 def test_save_wave_missing_audio_param(self, missing_param):
amine@132 601 with self.assertRaises(AudioParameterError):
amine@132 602 params = AUDIO_PARAMS_SHORT.copy()
amine@132 603 del params[missing_param]
amine@132 604 srate, swidth, channels, _ = _get_audio_parameters(params)
amine@136 605 _save_wave(b"\0\0", "audio", srate, swidth, channels)
amine@133 606
amine@141 607 def test_save_with_pydub(self):
amine@141 608 with patch("auditok.io.AudioSegment.export") as export:
amine@142 609 tmpdir = TemporaryDirectory()
amine@142 610 filename = os.path.join(tmpdir.name, "audio.ogg")
amine@142 611 _save_with_pydub(b"\0\0", filename, "ogg", 16000, 2, 1)
amine@141 612 self.assertTrue(export.called)
amine@142 613 tmpdir.cleanup()
amine@141 614
amine@133 615 @genty_dataset(
amine@133 616 raw_with_audio_format=("audio", "raw"),
amine@133 617 raw_with_extension=("audio.raw", None),
amine@133 618 raw_with_audio_format_and_extension=("audio.mp3", "raw"),
amine@133 619 raw_no_audio_format_nor_extension=("audio", None),
amine@133 620 )
amine@133 621 def test_to_file_raw(self, filename, audio_format):
amine@133 622 exp_filename = "tests/data/test_16KHZ_mono_400Hz.raw"
amine@133 623 tmpdir = TemporaryDirectory()
amine@133 624 filename = os.path.join(tmpdir.name, filename)
amine@133 625 data = _array_to_bytes(PURE_TONE_DICT[400])
amine@135 626 to_file(data, filename, audio_format=audio_format)
amine@133 627 self.assertTrue(filecmp.cmp(filename, exp_filename, shallow=False))
amine@133 628 tmpdir.cleanup()
amine@134 629
amine@134 630 @genty_dataset(
amine@134 631 wav_with_audio_format=("audio", "wav"),
amine@134 632 wav_with_extension=("audio.wav", None),
amine@134 633 wav_with_audio_format_and_extension=("audio.mp3", "wav"),
amine@134 634 wave_with_audio_format=("audio", "wave"),
amine@134 635 wave_with_extension=("audio.wave", None),
amine@134 636 wave_with_audio_format_and_extension=("audio.mp3", "wave"),
amine@134 637 )
amine@135 638 def test_to_file_wave(self, filename, audio_format):
amine@134 639 exp_filename = "tests/data/test_16KHZ_mono_400Hz.wav"
amine@134 640 tmpdir = TemporaryDirectory()
amine@134 641 filename = os.path.join(tmpdir.name, filename)
amine@134 642 data = _array_to_bytes(PURE_TONE_DICT[400])
amine@135 643 to_file(
amine@135 644 data,
amine@135 645 filename,
amine@135 646 audio_format=audio_format,
amine@135 647 sampling_rate=16000,
amine@135 648 sample_width=2,
amine@135 649 channels=1,
amine@134 650 )
amine@134 651 self.assertTrue(filecmp.cmp(filename, exp_filename, shallow=False))
amine@134 652 tmpdir.cleanup()
amine@138 653
amine@138 654 @genty_dataset(
amine@138 655 missing_sampling_rate=("sr",),
amine@138 656 missing_sample_width=("sw",),
amine@138 657 missing_channels=("ch",),
amine@138 658 )
amine@138 659 def test_to_file_missing_audio_param(self, missing_param):
amine@138 660 params = AUDIO_PARAMS_SHORT.copy()
amine@138 661 del params[missing_param]
amine@138 662 with self.assertRaises(AudioParameterError):
amine@138 663 to_file(b"\0\0", "audio", audio_format="wav", **params)
amine@138 664 with self.assertRaises(AudioParameterError):
amine@138 665 to_file(b"\0\0", "audio", audio_format="mp3", **params)
amine@139 666
amine@139 667 def test_to_file_no_pydub(self):
amine@139 668 with patch("auditok.io._WITH_PYDUB", False):
amine@139 669 with self.assertRaises(AudioIOError):
amine@139 670 to_file("audio", b"", "mp3")
amine@140 671
amine@140 672 @patch("auditok.io._WITH_PYDUB", True)
amine@140 673 @genty_dataset(
amine@140 674 ogg_with_extension=("audio.ogg", None),
amine@140 675 ogg_with_audio_format=("audio", "ogg"),
amine@140 676 ogg_format_with_wrong_extension=("audio.wav", "ogg"),
amine@140 677 )
amine@140 678 def test_to_file_compressed(self, filename, audio_format, *mocks):
amine@140 679 with patch("auditok.io.AudioSegment.export") as export:
amine@142 680 tmpdir = TemporaryDirectory()
amine@142 681 filename = os.path.join(tmpdir.name, filename)
amine@140 682 to_file(b"\0\0", filename, audio_format, **AUDIO_PARAMS_SHORT)
amine@140 683 self.assertTrue(export.called)
amine@142 684 tmpdir.cleanup()