|
8 | 8 | from io import BytesIO |
9 | 9 |
|
10 | 10 | import pytest |
11 | | -from ...testing import data_path, clear_and_catch_warnings, assert_arr_dict_equal |
| 11 | +from ...testing import data_path, clear_and_catch_warnings, assert_arr_dict_equal, error_warnings |
12 | 12 | from numpy.testing import assert_array_equal |
13 | 13 |
|
14 | 14 | from .test_tractogram import assert_tractogram_equal |
@@ -87,19 +87,19 @@ class TestTRK(unittest.TestCase): |
87 | 87 | def test_load_empty_file(self): |
88 | 88 | for lazy_load in [False, True]: |
89 | 89 | trk = TrkFile.load(DATA['empty_trk_fname'], lazy_load=lazy_load) |
90 | | - with pytest.warns(Warning if lazy_load else None): |
| 90 | + with pytest.warns(Warning) if lazy_load else error_warnings(): |
91 | 91 | assert_tractogram_equal(trk.tractogram, DATA['empty_tractogram']) |
92 | 92 |
|
93 | 93 | def test_load_simple_file(self): |
94 | 94 | for lazy_load in [False, True]: |
95 | 95 | trk = TrkFile.load(DATA['simple_trk_fname'], lazy_load=lazy_load) |
96 | | - with pytest.warns(Warning if lazy_load else None): |
| 96 | + with pytest.warns(Warning) if lazy_load else error_warnings(): |
97 | 97 | assert_tractogram_equal(trk.tractogram, DATA['simple_tractogram']) |
98 | 98 |
|
99 | 99 | def test_load_complex_file(self): |
100 | 100 | for lazy_load in [False, True]: |
101 | 101 | trk = TrkFile.load(DATA['complex_trk_fname'], lazy_load=lazy_load) |
102 | | - with pytest.warns(Warning if lazy_load else None): |
| 102 | + with pytest.warns(Warning) if lazy_load else error_warnings(): |
103 | 103 | assert_tractogram_equal(trk.tractogram, DATA['complex_tractogram']) |
104 | 104 |
|
105 | 105 | def trk_with_bytes(self, trk_key='simple_trk_fname', endian='<'): |
@@ -199,7 +199,7 @@ def test_load_complex_file_in_big_endian(self): |
199 | 199 | for lazy_load in [False, True]: |
200 | 200 | trk = TrkFile.load(DATA['complex_trk_big_endian_fname'], |
201 | 201 | lazy_load=lazy_load) |
202 | | - with pytest.warns(Warning if lazy_load else None): |
| 202 | + with pytest.warns(Warning) if lazy_load else error_warnings(): |
203 | 203 | assert_tractogram_equal(trk.tractogram, DATA['complex_tractogram']) |
204 | 204 |
|
205 | 205 | def test_tractogram_file_properties(self): |
|
0 commit comments