TST: Add ma.median tests for valid axis.

charris · charris · commit ad5b13a585cb · 2016-09-05T16:36:14.000-06:00
diff --git a/numpy/ma/tests/test_extras.py b/numpy/ma/tests/test_extras.py
@@ -10,6 +10,7 @@
 from __future__ import division, absolute_import, print_function
 
 import warnings
+import itertools
 
 import numpy as np
 from numpy.testing import (
@@ -684,6 +685,37 @@ def test_docstring_examples(self):
         assert_equal(ma_x.shape, (2,), "shape mismatch")
         assert_(type(ma_x) is MaskedArray)
 
+    def test_axis_argument_errors(self):
+        msg = "mask = %s, ndim = %s, axis = %s, overwrite_input = %s"
+        for ndmin in range(5):
+            for mask in [False, True]:
+                x = array(1, ndmin=ndmin, mask=mask)
+
+                # Valid axis values should not raise exception
+                args = itertools.product(range(-ndmin, ndmin), [False, True])
+                for axis, over in args:
+                    try:
+                        np.ma.median(x, axis=axis, overwrite_input=over)
+                    except:
+                        raise AssertionError(msg % (mask, ndmin, axis, over))
+
+                # Invalid axis values should raise exception
+                args = itertools.product([-(ndmin + 1), ndmin], [False, True])
+                for axis, over in args:
+                    try:
+                        np.ma.median(x, axis=axis, overwrite_input=over)
+                    except IndexError:
+                        pass
+                    else:
+                        raise AssertionError(msg % (mask, ndmin, axis, over))
+
+    def test_masked_0d(self):
+        # Check values
+        x = array(1, mask=False)
+        assert_equal(np.ma.median(x), 1)
+        x = array(1, mask=True)
+        assert_equal(np.ma.median(x), np.ma.masked)
+
     def test_masked_1d(self):
         x = array(np.arange(5), mask=True)
         assert_equal(np.ma.median(x), np.ma.masked)
