TST: Don't cast in test_cython_agg_nullable_int (pandas-dev#54248)

tpackard1 · Jul 26, 2023 · 317290a · 317290a
1 parent c9a8f95
commit 317290a
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/pandas/tests/groupby/aggregate/test_cython.py b/pandas/tests/groupby/aggregate/test_cython.py
@@ -338,13 +338,11 @@ def test_cython_agg_nullable_int(op_name):
     result = getattr(df.groupby("A")["B"], op_name)()
     df2 = df.assign(B=df["B"].astype("float64"))
     expected = getattr(df2.groupby("A")["B"], op_name)()
-
-    if op_name != "count":
-        # the result is not yet consistently using Int64/Float64 dtype,
-        # so for now just checking the values by casting to float
-        result = result.astype("float64")
+    if op_name in ("mean", "median"):
+        convert_integer = False
     else:
-        result = result.astype("int64")
+        convert_integer = True
+    expected = expected.convert_dtypes(convert_integer=convert_integer)
     tm.assert_series_equal(result, expected)