TST: group by - apply Key Error (#44806)

* Test for Group By - Apply Key Error

* Paramaters change

Author: brunoacpcosta

pandas/tests/groupby/test_apply.py

@@ -1177,3 +1177,25 @@ def test_apply_empty_string_nan_coerce_bug():
         index=MultiIndex.from_tuples([(1, ""), (2, "")], names=["a", "b"]),
     )
     tm.assert_frame_equal(result, expected)
+
+
+@pytest.mark.parametrize("index_values", [[1, 2, 3], [1.0, 2.0, 3.0]])
+def test_apply_index_key_error_bug(index_values):
+    # GH 44310
+    result = DataFrame(
+        {
+            "a": ["aa", "a2", "a3"],
+            "b": [1, 2, 3],
+        },
+        index=Index(index_values),
+    )
+    expected = DataFrame(
+        {
+            "b_mean": [2.0, 3.0, 1.0],
+        },
+        index=Index(["a2", "a3", "aa"], name="a"),
+    )
+    result = result.groupby("a").apply(
+        lambda df: Series([df["b"].mean()], index=["b_mean"])
+    )
+    tm.assert_frame_equal(result, expected)