Repository: pandas-dev/pandas
File: pandas/tests/arrays/test_integer.py (Python)

Code before:
    result = integer_array(values, dtype=to_dtype)
    assert result.dtype == result_dtype()
    expected = integer_array(values, dtype=result_dtype())
    tm.assert_extension_array_equal(result, expected)


def test_cross_type_arithmetic():

    df = pd.DataFrame({'A': pd.Series([1, 2, np.nan], dtype='Int64'),
                       'B': pd.Series([1, np.nan, 3], dtype='UInt8'),
                       'C': [1, 2, 3]})

    result = df.A + df.C
    expected = pd.Series([2, 4, np.nan], dtype='Int64')
    tm.assert_series_equal(result, expected)

    result = (df.A + df.C) * 3 == 12
    expected = pd.Series([False, True, False])
    tm.assert_series_equal(result, expected)

    result = df.A + df.B
    expected = pd.Series([2, np.nan, np.nan], dtype='Int64')
    tm.assert_series_equal(result, expected)


def test_arith_extension_array_values():
    # GH 22478
    s = pd.Series([1, 2, 3], dtype='Int64')
    result = s + s.values
    expected = pd.Series([2, 4, 6], dtype='Int64')
    tm.assert_series_equal(result, expected)


def test_groupby_mean_included():
    df = pd.DataFrame({
        "A": ['a', 'b', 'b'],
        "B": [1, None, 3],
        "C": integer_array([1, None, 3], dtype='Int64'),
    })

    result = df.groupby("A").sum()
    # TODO(#22346): preserve Int64 dtype
    expected = pd.DataFrame({
        "B": np.array([1.0, 3.0]),
        "C": np.array([1, 3], dtype="int64")
    }, index=pd.Index(['a', 'b'], name='A'))
    tm.assert_frame_equal(result, expected)


def test_astype_nansafe():
    # https://github.com/pandas-dev/pandas/pull/22343

Review comment (bug):
Can you move the test to `pandas/tests/extension/base/ops.py`, as this is not a specific failure for IntegerArray, it is relevant for all extension types.

Code after:
    result = integer_array(values, dtype=to_dtype)
    assert result.dtype == result_dtype()
    expected = integer_array(values, dtype=result_dtype())
    tm.assert_extension_array_equal(result, expected)


def test_cross_type_arithmetic():

    df = pd.DataFrame({'A': pd.Series([1, 2, np.nan], dtype='Int64'),
                       'B': pd.Series([1, np.nan, 3], dtype='UInt8'),
                       'C': [1, 2, 3]})

    result = df.A + df.C
    expected = pd.Series([2, 4, np.nan], dtype='Int64')
    tm.assert_series_equal(result, expected)

    result = (df.A + df.C) * 3 == 12
    expected = pd.Series([False, True, False])
    tm.assert_series_equal(result, expected)

    result = df.A + df.B
    expected = pd.Series([2, np.nan, np.nan], dtype='Int64')
    tm.assert_series_equal(result, expected)


def test_groupby_mean_included():
    df = pd.DataFrame({
        "A": ['a', 'b', 'b'],
        "B": [1, None, 3],
        "C": integer_array([1, None, 3], dtype='Int64'),
    })

    result = df.groupby("A").sum()
    # TODO(#22346): preserve Int64 dtype
    expected = pd.DataFrame({
        "B": np.array([1.0, 3.0]),
        "C": np.array([1, 3], dtype="int64")
    }, index=pd.Index(['a', 'b'], name='A'))
    tm.assert_frame_equal(result, expected)


def test_astype_nansafe():
    # https://github.com/pandas-dev/pandas/pull/22343
    arr = integer_array([np.nan, 1, 2], dtype="Int8")

    with tm.assert_raises_regex(
            ValueError, 'cannot convert float NaN to integer'):
        arr.astype('uint32')


# TODO(jreback) - these need testing / are broken