TST: add multi_index set tests on non-error cases

numpy/core/tests/test_nditer.py

@@ -185,6 +185,29 @@ def test_iter_c_or_f_order():
                 assert_equal([x for x in i],
                                     aview.swapaxes(0, 1).ravel(order='A'))
 
+def test_nditer_multi_index_set():
+    # Test the multi_index set
+    a = np.arange(6).reshape(2, 3)
+    it = np.nditer(a, flags=['multi_index'])
+
+    # Removes the iteration on two first elements of a[0]
+    it.multi_index = (0, 2,)
+
+    assert_equal([i for i in it], [2, 3, 4, 5])
+
+@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
+def test_nditer_multi_index_set_refcount():
+    # Test if the reference count on inde variable is decreased
+
+    index = 0
+    i = np.nditer(np.array([111, 222, 333, 444]), flags=['multi_index'])
+
+    start_count = sys.getrefcount(index)
+    i.multi_index = (index,)
+    end_count = sys.getrefcount(index)
+
+    assert_equal(start_count, end_count)
+
 def test_iter_best_order_multi_index_1d():
     # The multi-indices should be correct with any reordering