TST: GroupBy(..., as_index=True).agg() drops index (pandas-dev#33098)

Author: phofl

pandas/tests/groupby/aggregate/test_aggregate.py

@@ -809,7 +809,17 @@ def test_aggregate_mixed_types():
     tm.assert_frame_equal(result, expected)
 
 
-@pytest.mark.xfail(reason="Not implemented.")
+@pytest.mark.parametrize("func", ["min", "max"])
+def test_aggregate_categorical_lost_index(func: str):
+    # GH: 28641 groupby drops index, when grouping over categorical column with min/max
+    ds = pd.Series(["b"], dtype="category").cat.as_ordered()
+    df = pd.DataFrame({"A": [1997], "B": ds})
+    result = df.groupby("A").agg({"B": func})
+    expected = pd.DataFrame({"B": ["b"]}, index=pd.Index([1997], name="A"))
+    tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.xfail(reason="Not implemented;see GH 31256")
 def test_aggregate_udf_na_extension_type():
     # https://github.com/pandas-dev/pandas/pull/31359
     # This is currently failing to cast back to Int64Dtype.

pandas/tests/groupby/test_categorical.py

@@ -1388,6 +1388,19 @@ def test_groupby_agg_non_numeric():
     tm.assert_frame_equal(result, expected)
 
 
+@pytest.mark.parametrize("func", ["first", "last"])
+def test_groupy_first_returned_categorical_instead_of_dataframe(func):
+    # GH 28641: groupby drops index, when grouping over categorical column with
+    # first/last. Renamed Categorical instead of DataFrame previously.
+    df = pd.DataFrame(
+        {"A": [1997], "B": pd.Series(["b"], dtype="category").cat.as_ordered()}
+    )
+    df_grouped = df.groupby("A")["B"]
+    result = getattr(df_grouped, func)()
+    expected = pd.Series(["b"], index=pd.Index([1997], name="A"), name="B")
+    tm.assert_series_equal(result, expected)
+
+
 def test_read_only_category_no_sort():
     # GH33410
     cats = np.array([1, 2])