diff --git a/lib/matplotlib/tests/test_mlab.py b/lib/matplotlib/tests/test_mlab.py index 796b220f495c..fda6c955727f 100644 --- a/lib/matplotlib/tests/test_mlab.py +++ b/lib/matplotlib/tests/test_mlab.py @@ -2,8 +2,9 @@ unicode_literals) import six - +import warnings import tempfile +from functools import wraps from numpy.testing import assert_allclose, assert_array_equal import numpy.ma.testutils as matest @@ -1491,6 +1492,32 @@ def check_maxfreq(self, spec, fsp, fstims): del fstimst[-1] spect[maxind-5:maxind+5] = 0 + class catch_specgram_warnings(object): + # A decorator to catch warnings thrown by specgram + def __init__(self, n_specgram_calls): + self.n_specgram_calls = n_specgram_calls + + def __call__(self, test_function): + @wraps(test_function) + def wrapped(self_obj): + # Decide how many warnings will be thrown + n_warnings = 0 + if self_obj.NFFT_specgram is None: + if 256 >= len(self_obj.y): + n_warnings = 1 + elif self_obj.NFFT_specgram >= len(self_obj.y): + n_warnings = 1 + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter('always') + expected_warnings = self.n_specgram_calls * n_warnings + # Run test + test_function(self_obj) + assert len(w) == expected_warnings,\ + 'warnings:%s, expected:%s'%(len(w), expected_warnings) + + return wrapped + def test_spectral_helper_raises_complex_same_data(self): # test that mode 'complex' cannot be used if x is not y assert_raises(ValueError, mlab._spectral_helper, @@ -1936,6 +1963,7 @@ def test_phase_spectrum(self): assert_allclose(fsp, freqs, atol=1e-06) assert_equal(spec.shape, freqs.shape) + @catch_specgram_warnings(1) def test_specgram_auto(self): freqs = self.freqs_specgram spec, fsp, t = mlab.specgram(x=self.y, @@ -1945,7 +1973,6 @@ def test_specgram_auto(self): pad_to=self.pad_to_specgram, sides=self.sides) specm = np.mean(spec, axis=1) - assert_allclose(fsp, freqs, atol=1e-06) assert_allclose(t, self.t_specgram, atol=1e-06) @@ -1959,6 +1986,7 @@ def test_specgram_auto(self): atol=1e-02) self.check_freqs(specm, freqs, fsp, self.fstims) + @catch_specgram_warnings(1) def test_specgram_default(self): freqs = self.freqs_specgram spec, fsp, t = mlab.specgram(x=self.y, @@ -1983,6 +2011,7 @@ def test_specgram_default(self): atol=1e-02) self.check_freqs(specm, freqs, fsp, self.fstims) + @catch_specgram_warnings(1) def test_specgram_psd(self): freqs = self.freqs_specgram spec, fsp, t = mlab.specgram(x=self.y, @@ -2006,6 +2035,7 @@ def test_specgram_psd(self): atol=1e-02) self.check_freqs(specm, freqs, fsp, self.fstims) + @catch_specgram_warnings(1) def test_specgram_complex(self): freqs = self.freqs_specgram spec, fsp, t = mlab.specgram(x=self.y, @@ -2024,6 +2054,7 @@ def test_specgram_complex(self): self.check_freqs(specm, freqs, fsp, self.fstims) + @catch_specgram_warnings(1) def test_specgram_magnitude(self): freqs = self.freqs_specgram spec, fsp, t = mlab.specgram(x=self.y, @@ -2046,6 +2077,7 @@ def test_specgram_magnitude(self): atol=1e-02) self.check_freqs(specm, freqs, fsp, self.fstims) + @catch_specgram_warnings(1) def test_specgram_angle(self): freqs = self.freqs_specgram spec, fsp, t = mlab.specgram(x=self.y, @@ -2062,6 +2094,7 @@ def test_specgram_angle(self): assert_equal(spec.shape[0], freqs.shape[0]) assert_equal(spec.shape[1], self.t_specgram.shape[0]) + @catch_specgram_warnings(1) def test_specgram_phase(self): freqs = self.freqs_specgram spec, fsp, t = mlab.specgram(x=self.y, @@ -2096,6 +2129,7 @@ def test_psd_csd_equal(self): assert_array_equal(Pxx, Pxy) assert_array_equal(freqsxx, freqsxy) + @catch_specgram_warnings(2) def test_specgram_auto_default_equal(self): '''test that mlab.specgram without mode and with mode 'default' and 'psd' are all the same''' @@ -2117,6 +2151,7 @@ def test_specgram_auto_default_equal(self): assert_array_equal(freqspeca, freqspecb) assert_array_equal(ta, tb) + @catch_specgram_warnings(2) def test_specgram_auto_psd_equal(self): '''test that mlab.specgram without mode and with mode 'default' and 'psd' are all the same''' @@ -2138,6 +2173,7 @@ def test_specgram_auto_psd_equal(self): assert_array_equal(freqspeca, freqspecc) assert_array_equal(ta, tc) + @catch_specgram_warnings(2) def test_specgram_complex_mag_equivalent(self): freqs = self.freqs_specgram specc, freqspecc, tc = mlab.specgram(x=self.y, @@ -2159,6 +2195,7 @@ def test_specgram_complex_mag_equivalent(self): assert_array_equal(tc, tm) assert_allclose(np.abs(specc), specm, atol=1e-06) + @catch_specgram_warnings(2) def test_specgram_complex_angle_equivalent(self): freqs = self.freqs_specgram specc, freqspecc, tc = mlab.specgram(x=self.y, @@ -2180,6 +2217,7 @@ def test_specgram_complex_angle_equivalent(self): assert_array_equal(tc, ta) assert_allclose(np.angle(specc), speca, atol=1e-06) + @catch_specgram_warnings(2) def test_specgram_complex_phase_equivalent(self): freqs = self.freqs_specgram specc, freqspecc, tc = mlab.specgram(x=self.y, @@ -2202,6 +2240,7 @@ def test_specgram_complex_phase_equivalent(self): assert_allclose(np.unwrap(np.angle(specc), axis=0), specp, atol=1e-06) + @catch_specgram_warnings(2) def test_specgram_angle_phase_equivalent(self): freqs = self.freqs_specgram speca, freqspeca, ta = mlab.specgram(x=self.y,