TEST: Simplify lazy_load conditional warning check

Author: effigies

nibabel/streamlines/tests/test_streamlines.py

@@ -142,10 +142,9 @@ def test_load_empty_file(self):
                 else:
                     assert type(tfile.tractogram), LazyTractogram
 
-                with pytest.warns(None) as w:
+                with pytest.warns(Warning if lazy_load else None):
                     assert_tractogram_equal(tfile.tractogram,
                                             DATA['empty_tractogram'])
-                assert len(w) == lazy_load
 
     def test_load_simple_file(self):
         for lazy_load in [False, True]:
@@ -159,10 +158,9 @@ def test_load_simple_file(self):
                 else:
                     assert type(tfile.tractogram), LazyTractogram
 
-                with pytest.warns(None) as w:
+                with pytest.warns(Warning if lazy_load else None):
                     assert_tractogram_equal(tfile.tractogram,
                                             DATA['simple_tractogram'])
-                assert len(w) == lazy_load
 
     def test_load_complex_file(self):
         for lazy_load in [False, True]:
@@ -186,10 +184,9 @@ def test_load_complex_file(self):
                     data = DATA['data_per_streamline']
                     tractogram.data_per_streamline = data
 
-                with pytest.warns(None) as w:
+                with pytest.warns(Warning if lazy_load else None):
                     assert_tractogram_equal(tfile.tractogram,
                                             tractogram)
-                assert len(w) == lazy_load
 
     def test_save_tractogram_file(self):
         tractogram = Tractogram(DATA['streamlines'],

nibabel/streamlines/tests/test_tck.py

@@ -47,16 +47,14 @@ class TestTCK(unittest.TestCase):
     def test_load_empty_file(self):
         for lazy_load in [False, True]:
             tck = TckFile.load(DATA['empty_tck_fname'], lazy_load=lazy_load)
-            with pytest.warns(None) as w:
+            with pytest.warns(Warning if lazy_load else None):
                 assert_tractogram_equal(tck.tractogram, DATA['empty_tractogram'])
-            assert len(w) == lazy_load
 
     def test_load_simple_file(self):
         for lazy_load in [False, True]:
             tck = TckFile.load(DATA['simple_tck_fname'], lazy_load=lazy_load)
-            with pytest.warns(None) as w:
+            with pytest.warns(Warning if lazy_load else None):
                 assert_tractogram_equal(tck.tractogram, DATA['simple_tractogram'])
-            assert len(w) == lazy_load
 
         # Force TCK loading to use buffering.
         buffer_size = 1. / 1024**2  # 1 bytes
@@ -90,9 +88,8 @@ def test_load_simple_file_in_big_endian(self):
         for lazy_load in [False, True]:
             tck = TckFile.load(DATA['simple_tck_big_endian_fname'],
                                lazy_load=lazy_load)
-            with pytest.warns(None) as w:
+            with pytest.warns(Warning if lazy_load else None):
                 assert_tractogram_equal(tck.tractogram, DATA['simple_tractogram'])
-            assert len(w) == lazy_load
             assert tck.header['datatype'] == 'Float32BE'
 
     def test_load_file_with_wrong_information(self):

nibabel/streamlines/tests/test_trk.py

@@ -87,23 +87,20 @@ class TestTRK(unittest.TestCase):
     def test_load_empty_file(self):
         for lazy_load in [False, True]:
             trk = TrkFile.load(DATA['empty_trk_fname'], lazy_load=lazy_load)
-            with pytest.warns(None) as w:
+            with pytest.warns(Warning if lazy_load else None):
                 assert_tractogram_equal(trk.tractogram, DATA['empty_tractogram'])
-            assert len(w) == lazy_load
 
     def test_load_simple_file(self):
         for lazy_load in [False, True]:
             trk = TrkFile.load(DATA['simple_trk_fname'], lazy_load=lazy_load)
-            with pytest.warns(None) as w:
+            with pytest.warns(Warning if lazy_load else None):
                 assert_tractogram_equal(trk.tractogram, DATA['simple_tractogram'])
-            assert len(w) == lazy_load
 
     def test_load_complex_file(self):
         for lazy_load in [False, True]:
             trk = TrkFile.load(DATA['complex_trk_fname'], lazy_load=lazy_load)
-            with pytest.warns(None) as w:
+            with pytest.warns(Warning if lazy_load else None):
                 assert_tractogram_equal(trk.tractogram, DATA['complex_tractogram'])
-            assert len(w) == lazy_load
 
     def trk_with_bytes(self, trk_key='simple_trk_fname', endian='<'):
         """ Return example trk file bytes and struct view onto bytes """
@@ -202,9 +199,8 @@ def test_load_complex_file_in_big_endian(self):
         for lazy_load in [False, True]:
             trk = TrkFile.load(DATA['complex_trk_big_endian_fname'],
                                lazy_load=lazy_load)
-            with pytest.warns(None) as w:
+            with pytest.warns(Warning if lazy_load else None):
                 assert_tractogram_equal(trk.tractogram, DATA['complex_tractogram'])
-            assert len(w) == lazy_load
 
     def test_tractogram_file_properties(self):
         trk = TrkFile.load(DATA['simple_trk_fname'])