Update test scripts

Author: Daniel Saxton

pandas/tests/arrays/test_integer.py

@@ -1074,15 +1074,21 @@ def test_cut(bins, right, include_lowest):
     tm.assert_categorical_equal(result, expected)
 
 
-@pytest.mark.parametrize("q", [2, 5, 10])
-def test_qcut_nullable_integer(q, any_nullable_int_dtype):
-    arr = pd.array(np.arange(100), dtype=any_nullable_int_dtype)
-    arr[::2] = pd.NA
+def test_array_setitem_nullable_boolean_mask():
+    # GH 31446
+    ser = pd.Series([1, 2], dtype="Int64")
+    result = ser.where(ser > 1)
+    expected = pd.Series([pd.NA, 2], dtype="Int64")
+    tm.assert_series_equal(result, expected)
 
-    result = pd.qcut(arr, q)
-    expected = pd.qcut(arr.astype(float), q)
 
-    tm.assert_categorical_equal(result, expected)
+def test_array_setitem():
+    # GH 31446
+    arr = pd.Series([1, 2], dtype="Int64").array
+    arr[arr > 1] = 1
+
+    expected = pd.array([1, 1], dtype="Int64")
+    tm.assert_extension_array_equal(arr, expected)
 
 
 # TODO(jreback) - these need testing / are broken

pandas/tests/reshape/test_qcut.py

@@ -4,6 +4,7 @@
 import pytest
 
 from pandas import (
+    NA,
     Categorical,
     DatetimeIndex,
     Interval,
@@ -12,6 +13,7 @@
     Series,
     TimedeltaIndex,
     Timestamp,
+    array,
     cut,
     date_range,
     isna,
@@ -286,3 +288,14 @@ def test_qcut_bool_coercion_to_int(bins, box, compare):
     expected = qcut(data_expected, bins, duplicates="drop")
     result = qcut(data_result, bins, duplicates="drop")
     compare(result, expected)
+
+
+@pytest.mark.parametrize("q", [2, 5, 10])
+def test_qcut_nullable_integer(q, any_nullable_int_dtype):
+    arr = array(np.arange(100), dtype=any_nullable_int_dtype)
+    arr[::2] = NA
+
+    result = qcut(arr, q)
+    expected = qcut(arr.astype(float), q)
+
+    tm.assert_categorical_equal(result, expected)