updated tests

Author: bobhaffner

pandas/tests/reshape/test_pivot.py

@@ -1109,18 +1109,46 @@ def test_pivot_margins_name_unicode(self):
         expected = pd.DataFrame(index=index)
         tm.assert_frame_equal(table, expected)
 
-    def test_pivot_func_strings(self):
+    def test_pivot_string_as_func(self):
         # GH #18713
-        f = lambda func: pivot_table(self.data, values=['D', 'E'],
-                                     index=['A', 'B'], columns='C',
-                                     aggfunc=func)
-        result = f('sum')
-        expected = f(np.sum)
+        data = DataFrame({'A': ['foo', 'foo', 'foo', 'foo', 'bar', 'bar',
+                                'bar', 'bar', 'foo', 'foo', 'foo'],
+                          'B': ['one', 'one', 'one', 'two', 'one', 'one',
+                                'one', 'two', 'two', 'two', 'one'],
+                          'C': range(11)})
+
+        result = pivot_table(data, index='A', columns='B', aggfunc='sum')
+        mi = MultiIndex(levels=[['C'], ['one', 'two']],
+                        labels=[[0, 0], [0, 1]], names=[None, 'B'])
+        expected = DataFrame({('C', 'one'): {'bar': 15, 'foo': 13},
+                              ('C', 'two'): {'bar': 7, 'foo': 20}},
+                             columns=mi).rename_axis('A')
         tm.assert_frame_equal(result, expected)
-        result = f(['mean', 'std'])
-        means = f(np.mean)
-        stds = f(np.std)
-        expected = concat([means, stds], keys=['mean', 'std'], axis=1)
+
+        result = pivot_table(data, index='A', columns='B',
+                             aggfunc=['sum', 'mean'])
+        mi = MultiIndex(levels=[['sum', 'mean'], ['C'], ['one', 'two']],
+                        labels=[[0, 0, 1, 1], [0, 0, 0, 0], [0, 1, 0, 1]],
+                        names=[None, None, 'B'])
+        expected = DataFrame({('mean', 'C', 'one'): {'bar': 5.0, 'foo': 3.25},
+                              ('mean', 'C', 'two'): {'bar': 7.0,
+                                                     'foo': 6.666666666666667},
+                              ('sum', 'C', 'one'): {'bar': 15, 'foo': 13},
+                              ('sum', 'C', 'two'): {'bar': 7, 'foo': 20}},
+                             columns=mi).rename_axis('A')
+        tm.assert_frame_equal(result, expected)
+
+    funcs = [('sum', np.sum), ('mean', np.mean), ('std', np.std),
+             (['sum', 'mean'], [np.sum, np.mean]),
+             (['sum', 'std'], [np.sum, np.std]),
+             (['std', 'mean'], [np.std, np.mean])]
+
+    @pytest.mark.parametrize("f, f_numpy", funcs)
+    def test_pivot_string_func_vs_func(self, f, f_numpy):
+        # GH #18713
+        result = pivot_table(self.data, index='A', columns='B', aggfunc=f)
+        expected = pivot_table(self.data, index='A', columns='B',
+                               aggfunc=f_numpy)
         tm.assert_frame_equal(result, expected)