Merge pull request #12090 from pierreglaser/add-protocol-kwarg-to-array-dump

MNT Update pickling test by making them loop over all protocols

1 files changed, 5 insertions, 4 deletions

diff --git a/numpy/matrixlib/tests/test_masked_matrix.py b/numpy/matrixlib/tests/test_masked_matrix.py
index 5ed8044aa..7f84bb2c9 100644
--- a/numpy/matrixlib/tests/test_masked_matrix.py
+++ b/numpy/matrixlib/tests/test_masked_matrix.py
@@ -79,10 +79,11 @@ class TestMaskedMatrix(object):
     def test_pickling_subbaseclass(self):
         # Test pickling w/ a subclass of ndarray
         a = masked_array(np.matrix(list(range(10))), mask=[1, 0, 1, 0, 0] * 2)
-        a_pickled = pickle.loads(a.dumps())
-        assert_equal(a_pickled._mask, a._mask)
-        assert_equal(a_pickled, a)
-        assert_(isinstance(a_pickled._data, np.matrix))
+        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
+            a_pickled = pickle.loads(pickle.dumps(a, protocol=proto))
+            assert_equal(a_pickled._mask, a._mask)
+            assert_equal(a_pickled, a)
+            assert_(isinstance(a_pickled._data, np.matrix))
 
     def test_count_mean_with_matrix(self):
         m = masked_array(np.matrix([[1, 2], [3, 4]]), mask=np.zeros((2, 2)))