addtl tests

pandas-dev · May 31, 2018 · 5c70f0d · 5c70f0d
1 parent 1745fb5
commit 5c70f0d
Showing 1 changed file with 23 additions and 7 deletions.
diff --git a/pandas/tests/reshape/test_pivot.py b/pandas/tests/reshape/test_pivot.py
@@ -151,14 +151,30 @@ def test_pivot_with_non_observable_dropna(self, dropna):
              'B': range(5)})
 
         result = df.pivot_table(index='A', values='B', dropna=dropna)
-        expected = pd.DataFrame(
-            {'B': [2, 3]},
-            index=pd.Index(
-                pd.Categorical.from_codes([0, 1],
-                                          categories=['low', 'high'],
-                                          ordered=True),
-                name='A'))
+        expected_index = pd.CategoricalIndex(['low', 'high'],
+                                             categories=['low', 'high'],
+                                             ordered=True, name='A')
+
+        expected = pd.DataFrame({'B': [2, 3]},
+                                index=expected_index)
+
+        tm.assert_frame_equal(result, expected)
 
+        result = df.pivot_table(columns='A', values='B', dropna=dropna)
+        expected = pd.DataFrame([[2, 3]],
+                                index=pd.Index(['B']),
+                                columns=expected_index)
+        tm.assert_frame_equal(result, expected)
+
+        df['AA'] = df['A']
+        result = df.pivot_table(index='A',
+                                columns='AA',
+                                values='B',
+                                dropna=dropna)
+        expected = pd.DataFrame(
+            [[2, np.nan], [np.nan, 3]],
+            index=expected_index,
+            columns=expected_index.rename('AA'))
         tm.assert_frame_equal(result, expected)
 
     def test_pass_array(self):