Refactor groupby-resample tests via TestCase to remove duplicate code.

pooja-subramaniam · Jan 20, 2023 · d522606 · d522606
1 parent 645cf1f
commit d522606
Showing 1 changed file with 61 additions and 77 deletions.
diff --git a/pandas/tests/resample/test_resampler_grouper.py b/pandas/tests/resample/test_resampler_grouper.py
@@ -1,3 +1,4 @@
+import unittest
 from textwrap import dedent
 
 import numpy as np
@@ -538,80 +539,63 @@ def test_groupby_resample_size_all_index_same():
     tm.assert_series_equal(result, expected)
 
 
-def test_groupby_resample_on_index_with_list_of_keys():
-    # GH 50840
-    df = DataFrame(
-        data={
-            "group": [0, 0, 0, 0, 1, 1, 1, 1],
-            "val": [3, 1, 4, 1, 5, 9, 2, 6],
-        },
-        index=Series(
-            date_range(start="2016-01-01", periods=8),
-            name="date",
-        ),
-    )
-    result = df.groupby("group").resample("2D")[["val"]].mean()
-    expected = DataFrame(
-        data={
-            "val": [2.0, 2.5, 7.0, 4.0],
-        },
-        index=Index(
-            data=[
-                (0, Timestamp("2016-01-01")),
-                (0, Timestamp("2016-01-03")),
-                (1, Timestamp("2016-01-05")),
-                (1, Timestamp("2016-01-07")),
-            ],
-            name=("group", "date"),
-        ),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-def test_groupby_resample_on_index_with_list_of_keys_multi_columns():
-    # GH 50876
-    df = DataFrame(
-        data={
-            "group": [0, 0, 0, 0, 1, 1, 1, 1],
-            "first_val": [3, 1, 4, 1, 5, 9, 2, 6],
-            "second_val": [2, 7, 1, 8, 2, 8, 1, 8],
-            "third_val": [1, 4, 1, 4, 2, 1, 3, 5],
-        },
-        index=Series(
-            date_range(start="2016-01-01", periods=8),
-            name="date",
-        ),
-    )
-    result = df.groupby("group").resample("2D")[["first_val", "second_val"]].mean()
-    expected = DataFrame(
-        data={
-            "first_val": [2.0, 2.5, 7.0, 4.0],
-            "second_val": [4.5, 4.5, 5.0, 4.5],
-        },
-        index=Index(
-            data=[
-                (0, Timestamp("2016-01-01")),
-                (0, Timestamp("2016-01-03")),
-                (1, Timestamp("2016-01-05")),
-                (1, Timestamp("2016-01-07")),
-            ],
-            name=("group", "date"),
-        ),
-    )
-    tm.assert_frame_equal(result, expected)
-
-
-def test_groupby_resample_on_index_with_list_of_keys_missing_column():
-    # GH 50876
-    df = DataFrame(
-        data={
-            "group": [0, 0, 0, 0, 1, 1, 1, 1],
-            "val": [3, 1, 4, 1, 5, 9, 2, 6],
-        },
-        index=Series(
-            date_range(start="2016-01-01", periods=8),
-            name="date",
-        ),
-    )
-    with pytest.raises(KeyError, match="Columns not found"):
-        df.groupby("group").resample("2D")[["val_not_in_dataframe"]].mean()
+class TestGroupByResampleTimeIndex(unittest.TestCase):
+    """Test groupby resample with a time index where a list of columns is given."""
+    def setUp(self) -> None:
+        self.df = DataFrame(
+            data={
+                "group": [0, 0, 0, 0, 1, 1, 1, 1],
+                "first_val": [3, 1, 4, 1, 5, 9, 2, 6],
+                "second_val": [2, 7, 1, 8, 2, 8, 1, 8],
+                "third_val": [1, 4, 1, 4, 2, 1, 3, 5],
+            },
+            index=Series(
+                date_range(start="2016-01-01", periods=8),
+                name="date",
+            ),
+        )
+
+    def test_list_of_one_key(self):
+        # GH 50840
+        result = self.df.groupby("group").resample("2D")[["first_val"]].mean()
+        expected = DataFrame(
+            data={
+                "first_val": [2.0, 2.5, 7.0, 4.0],
+            },
+            index=Index(
+                data=[
+                    (0, Timestamp("2016-01-01")),
+                    (0, Timestamp("2016-01-03")),
+                    (1, Timestamp("2016-01-05")),
+                    (1, Timestamp("2016-01-07")),
+                ],
+                name=("group", "date"),
+            ),
+        )
+        tm.assert_frame_equal(result, expected)
+
+    def test_list_of_multiple_keys(self):
+        # GH 50876
+        result = self.df.groupby("group").resample("2D")[["first_val", "second_val"]].mean()
+        expected = DataFrame(
+            data={
+                "first_val": [2.0, 2.5, 7.0, 4.0],
+                "second_val": [4.5, 4.5, 5.0, 4.5],
+            },
+            index=Index(
+                data=[
+                    (0, Timestamp("2016-01-01")),
+                    (0, Timestamp("2016-01-03")),
+                    (1, Timestamp("2016-01-05")),
+                    (1, Timestamp("2016-01-07")),
+                ],
+                name=("group", "date"),
+            ),
+        )
+        tm.assert_frame_equal(result, expected)
+
+    def test_missing_key_raises_KeyError(self):
+        """Test a key that is not in the list of columns."""
+        # GH 50876
+        with pytest.raises(KeyError, match="Columns not found"):
+            self.df.groupby("group").resample("2D")[["val_not_in_dataframe"]].mean()