From fbe103c129dfcd43051714d08557363ce65f0e01 Mon Sep 17 00:00:00 2001
From: Richard Shadrach <rhshadrach@gmail.com>
Date: Sun, 8 Nov 2020 10:05:14 -0500
Subject: [PATCH] CLN: Simplify groupby head/tail tests

---
 pandas/tests/groupby/test_nth.py | 69 +++++++++++---------------------
 1 file changed, 23 insertions(+), 46 deletions(-)

diff --git a/pandas/tests/groupby/test_nth.py b/pandas/tests/groupby/test_nth.py
index df1d7819a1894..10394ea997775 100644
--- a/pandas/tests/groupby/test_nth.py
+++ b/pandas/tests/groupby/test_nth.py
@@ -487,53 +487,30 @@ def test_nth_multi_index_as_expected():
     tm.assert_frame_equal(result, expected)
 
 
-def test_groupby_head_tail():
+@pytest.mark.parametrize(
+    "op, n, expected_rows",
+    [
+        ("head", -1, []),
+        ("head", 0, []),
+        ("head", 1, [0, 2]),
+        ("head", 7, [0, 1, 2]),
+        ("tail", -1, []),
+        ("tail", 0, []),
+        ("tail", 1, [1, 2]),
+        ("tail", 7, [0, 1, 2]),
+    ],
+)
+@pytest.mark.parametrize("columns", [None, [], ["A"], ["B"], ["A", "B"]])
+@pytest.mark.parametrize("as_index", [True, False])
+def test_groupby_head_tail(op, n, expected_rows, columns, as_index):
     df = DataFrame([[1, 2], [1, 4], [5, 6]], columns=["A", "B"])
-    g_as = df.groupby("A", as_index=True)
-    g_not_as = df.groupby("A", as_index=False)
-
-    # as_index= False, much easier
-    tm.assert_frame_equal(df.loc[[0, 2]], g_not_as.head(1))
-    tm.assert_frame_equal(df.loc[[1, 2]], g_not_as.tail(1))
-
-    empty_not_as = DataFrame(columns=df.columns, index=Index([], dtype=df.index.dtype))
-    empty_not_as["A"] = empty_not_as["A"].astype(df.A.dtype)
-    empty_not_as["B"] = empty_not_as["B"].astype(df.B.dtype)
-    tm.assert_frame_equal(empty_not_as, g_not_as.head(0))
-    tm.assert_frame_equal(empty_not_as, g_not_as.tail(0))
-    tm.assert_frame_equal(empty_not_as, g_not_as.head(-1))
-    tm.assert_frame_equal(empty_not_as, g_not_as.tail(-1))
-
-    tm.assert_frame_equal(df, g_not_as.head(7))  # contains all
-    tm.assert_frame_equal(df, g_not_as.tail(7))
-
-    # as_index=True, (used to be different)
-    df_as = df
-
-    tm.assert_frame_equal(df_as.loc[[0, 2]], g_as.head(1))
-    tm.assert_frame_equal(df_as.loc[[1, 2]], g_as.tail(1))
-
-    empty_as = DataFrame(index=df_as.index[:0], columns=df.columns)
-    empty_as["A"] = empty_not_as["A"].astype(df.A.dtype)
-    empty_as["B"] = empty_not_as["B"].astype(df.B.dtype)
-    tm.assert_frame_equal(empty_as, g_as.head(0))
-    tm.assert_frame_equal(empty_as, g_as.tail(0))
-    tm.assert_frame_equal(empty_as, g_as.head(-1))
-    tm.assert_frame_equal(empty_as, g_as.tail(-1))
-
-    tm.assert_frame_equal(df_as, g_as.head(7))  # contains all
-    tm.assert_frame_equal(df_as, g_as.tail(7))
-
-    # test with selection
-    tm.assert_frame_equal(g_as[[]].head(1), df_as.loc[[0, 2], []])
-    tm.assert_frame_equal(g_as[["A"]].head(1), df_as.loc[[0, 2], ["A"]])
-    tm.assert_frame_equal(g_as[["B"]].head(1), df_as.loc[[0, 2], ["B"]])
-    tm.assert_frame_equal(g_as[["A", "B"]].head(1), df_as.loc[[0, 2]])
-
-    tm.assert_frame_equal(g_not_as[[]].head(1), df_as.loc[[0, 2], []])
-    tm.assert_frame_equal(g_not_as[["A"]].head(1), df_as.loc[[0, 2], ["A"]])
-    tm.assert_frame_equal(g_not_as[["B"]].head(1), df_as.loc[[0, 2], ["B"]])
-    tm.assert_frame_equal(g_not_as[["A", "B"]].head(1), df_as.loc[[0, 2]])
+    g = df.groupby("A", as_index=as_index)
+    expected = df.iloc[expected_rows]
+    if columns is not None:
+        g = g[columns]
+        expected = expected[columns]
+    result = getattr(g, op)(n)
+    tm.assert_frame_equal(result, expected)
 
 
 def test_group_selection_cache():