amine@347: import os
amine@348: import sys
amine@347: import unittest
amine@347: from unittest import TestCase
amine@347: from tempfile import TemporaryDirectory
amine@347: from genty import genty, genty_dataset
amine@347: import matplotlib
amine@348: 
amine@348: matplotlib.use("AGG")  # noqa E402
amine@347: import matplotlib.pyplot as plt
amine@347: from auditok.core import AudioRegion
amine@347: 
amine@348: if sys.version_info.minor <= 5:
amine@348:     PREFIX = "py34_py35/"
amine@348: else:
amine@348:     PREFIX = ""
amine@348: 
amine@347: matplotlib.rcParams["figure.figsize"] = (10, 4)
amine@347: 
amine@347: 
amine@347: @genty
amine@347: class TestPlotting(TestCase):
amine@347:     @genty_dataset(mono=(1,), stereo=(2,))
amine@347:     def test_region_plot(self, channels):
amine@347:         type_ = "mono" if channels == 1 else "stereo"
amine@347:         audio_filename = "tests/data/test_split_10HZ_{}.raw".format(type_)
amine@348:         image_filename = "tests/images/{}plot_{}_region.png".format(
amine@348:             PREFIX, type_
amine@348:         )
amine@347:         expected_image = plt.imread(image_filename)
amine@347:         with TemporaryDirectory() as tmpdir:
amine@347:             output_image_filename = os.path.join(tmpdir, "image.png")
amine@347:             region = AudioRegion.load(audio_filename, sr=10, sw=2, ch=channels)
amine@347:             region.plot(show=False, save_as=output_image_filename)
amine@347:             output_image = plt.imread(output_image_filename)
amine@347:         self.assertTrue((output_image == expected_image).all())
amine@347: 
amine@347:     @genty_dataset(
amine@347:         mono=(1,),
amine@347:         stereo_any=(2, "any"),
amine@347:         stereo_uc_0=(2, 0),
amine@347:         stereo_uc_1=(2, 1),
amine@347:         stereo_uc_mix=(2, "mix"),
amine@347:     )
amine@347:     def test_region_split_and_plot(self, channels, use_channel=None):
amine@347:         type_ = "mono" if channels == 1 else "stereo"
amine@347:         audio_filename = "tests/data/test_split_10HZ_{}.raw".format(type_)
amine@347:         if type_ == "mono":
amine@348:             fmt = "tests/images/{}split_and_plot_mono_region.png"
amine@347:         else:
amine@348:             fmt = "tests/images/{}split_and_plot_uc_{}_stereo_region.png"
amine@348:         image_filename = fmt.format(PREFIX, use_channel)
amine@347: 
amine@347:         expected_image = plt.imread(image_filename)
amine@347:         with TemporaryDirectory() as tmpdir:
amine@347:             output_image_filename = os.path.join(tmpdir, "image.png")
amine@347:             region = AudioRegion.load(audio_filename, sr=10, sw=2, ch=channels)
amine@347:             region.split_and_plot(
amine@347:                 aw=0.1,
amine@347:                 uc=use_channel,
amine@347:                 max_silence=0,
amine@347:                 show=False,
amine@347:                 save_as=output_image_filename,
amine@347:             )
amine@347:             output_image = plt.imread(output_image_filename)
amine@347:         self.assertTrue((output_image == expected_image).all())
amine@347: 
amine@347: 
amine@347: if __name__ == "__main__":
amine@347:     unittest.main()