Compare resulting DataFrame in tests

Author: BryanCutler

pandas/tests/extension/base/groupby.py

@@ -36,11 +36,18 @@ def test_groupby_extension_agg(self, as_index, data_for_grouping):
     def test_groupby_agg_extension(self, data_for_grouping):
         # GH#38980 groupby agg on extension type fails for non-numeric types
         df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping})
-        result = df.groupby("A").agg({"B": "first"}).B.array
 
-        expected = df["B"].iloc[[0, 2, 4, 7]].array
+        expected = df.iloc[[0, 2, 4, 7]]
+        expected = expected.set_index("A")
 
-        self.assert_extension_array_equal(result, expected)
+        result = df.groupby("A").agg({"B": "first"})
+        self.assert_frame_equal(result, expected)
+
+        result = df.groupby("A").agg("first")
+        self.assert_frame_equal(result, expected)
+
+        result = df.groupby("A").first()
+        self.assert_frame_equal(result, expected)
 
     def test_groupby_extension_no_sort(self, data_for_grouping):
         df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1, 4], "B": data_for_grouping})

pandas/tests/extension/test_boolean.py

@@ -294,11 +294,18 @@ def test_groupby_extension_agg(self, as_index, data_for_grouping):
     def test_groupby_agg_extension(self, data_for_grouping):
         # GH#38980 groupby agg on extension type fails for non-numeric types
         df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})
-        result = df.groupby("A").agg({"B": "first"}).B.array
 
-        expected = df["B"].iloc[[0, 2, 4]].array
+        expected = df.iloc[[0, 2, 4]]
+        expected = expected.set_index("A")
 
-        self.assert_extension_array_equal(result, expected)
+        result = df.groupby("A").agg({"B": "first"})
+        self.assert_frame_equal(result, expected)
+
+        result = df.groupby("A").agg("first")
+        self.assert_frame_equal(result, expected)
+
+        result = df.groupby("A").first()
+        self.assert_frame_equal(result, expected)
 
     def test_groupby_extension_no_sort(self, data_for_grouping):
         df = pd.DataFrame({"A": [1, 1, 2, 2, 3, 3, 1], "B": data_for_grouping})