[test] add a test using sink and source to check channel mixing
[aubio.git] / python / tests / test_source_channels.py
1 #! /usr/bin/env python
2
3 """A brute force test using `sink` to create and write samples to a stereo
4 file, then `source` to check the correct content is read from the files."""
5
6 import os.path
7 import unittest
8 import aubio
9 import numpy as np
10 from numpy.testing import assert_equal
11 from utils import get_tmp_sink_path
12
13 class aubio_source_test_case(unittest.TestCase):
14
15     def test_read_from_mono(self):
16         out = get_tmp_sink_path()
17         samplerate = 44100
18         hop_size = 256
19         blocks = 10
20         channels = 1
21         write_samples = np.ones([channels, hop_size], dtype=aubio.float_type)
22         write_samples *= .5
23         self.check_write_and_read(samplerate, channels, hop_size, blocks,
24                 write_samples)
25
26     def test_read_from_stereo(self):
27         out = get_tmp_sink_path()
28         samplerate = 44100
29         hop_size = 256
30         blocks = 10
31         channels = 1
32         write_samples = np.ones([channels, hop_size], dtype=aubio.float_type)
33         write_samples *= .5
34         self.check_write_and_read(samplerate, channels, hop_size, blocks,
35                 write_samples)
36
37     def test_read_from_half_stereo(self):
38         samplerate = 16000
39         channels = 2
40         hop_size = 512
41         blocks = 10
42         write_samples = np.ones([channels, hop_size], dtype=aubio.float_type)
43         write_samples *= .5
44         write_samples[1, :] = 0
45         self.check_write_and_read(samplerate, channels, hop_size, blocks,
46                 write_samples)
47
48     def test_read_from_cancelling_channels(self):
49         samplerate = 16000
50         channels = 2
51         hop_size = 512
52         blocks = 10
53         write_samples = np.ones([channels, hop_size], dtype=aubio.float_type)
54         write_samples *= .5
55         write_samples[1] *= -1
56         self.check_write_and_read(samplerate, channels, hop_size, blocks,
57                 write_samples)
58
59     def test_read_from_strange_three_channels(self):
60         samplerate = 8000
61         channels = 3
62         hop_size = 123
63         blocks = 10
64         write_samples = np.ones([channels, hop_size], dtype=aubio.float_type)
65         write_samples *= .5
66         write_samples[1, :] = 0
67         self.check_write_and_read(samplerate, channels, hop_size, blocks,
68                 write_samples)
69
70     def check_write_and_read(self, samplerate, channels,
71             hop_size, blocks, write_samples):
72         expected_mono = np.sum(write_samples, axis=0)/write_samples.shape[0]
73         out = get_tmp_sink_path()
74         snk = aubio.sink(out, samplerate, channels=channels)
75         for i in range(blocks):
76             snk.do_multi(write_samples, hop_size)
77         # close the sink before reading from it
78         snk.close()
79
80         src = aubio.source(out, samplerate, hop_size)
81         for i in range(blocks):
82             read_samples, read = src.do_multi()
83             assert_equal (read_samples, write_samples)
84             assert_equal (read, hop_size)
85
86         src.seek(0)
87         for i in range(blocks):
88             read_samples, read = src()
89             assert_equal (read, hop_size)
90             assert_equal (read_samples, expected_mono)
91
92 if __name__ == '__main__':
93     unittest.main()