test case dataframe assertion added

Author: 5hirish

pandas/core/groupby/groupby.py

@@ -28,7 +28,7 @@ class providing the base-class of operations.
 from pandas.core.config import option_context
 from pandas.core.dtypes.cast import maybe_downcast_to_dtype
 from pandas.core.dtypes.common import (
-    ensure_float, is_numeric_dtype, is_scalar, is_extension_array_dtype
+    ensure_float, is_extension_array_dtype, is_numeric_dtype, is_scalar
 )
 from pandas.core.dtypes.missing import isna, notna
 from pandas.core.frame import DataFrame

pandas/tests/arrays/test_integer.py

@@ -674,22 +674,11 @@ def test_reduce_to_float(op):
     # groupby
     result = getattr(df.groupby("A"), op)()
 
-    assert result.dtypes['B'].name == 'float64'
-    assert result.dtypes['C'].name == 'Int64'
-
-
-@pytest.mark.parametrize('op', ['sum'])
-def test_groupby_extension_array(op):
-    # GH23227
-    # groupby on an extension array should return the extension array type
-    df = pd.DataFrame({
-        'Int': pd.Series([1, 2, 3], dtype='Int64'),
-        'A': [1, 2, 1]
-    })
-
-    result = getattr(df.groupby('A').Int, op)()
-    assert result is not None
-    assert result.dtype.name == 'Int64'
+    expected = pd.DataFrame({
+        "B": np.array([1.0, 3.0]),
+        "C": integer_array([1, 3], dtype="Int64")
+    }, index=pd.Index(['a', 'b'], name='A'))
+    tm.assert_frame_equal(result, expected)
 
 
 def test_astype_nansafe():