Add tests for single level indexing with loc(axis=1)

Author: ryankarlos

pandas/tests/indexing/multiindex/test_slice.py

@@ -492,6 +492,29 @@ def test_loc_axis_arguments(self):
         with pytest.raises(ValueError):
             df.loc(axis="foo")[:, :, ["C1", "C3"]]
 
+    def test_loc_axis_single_level_indexer(self):
+
+        # test single level indexing on Multindex column dataframe
+        df1 = pd.DataFrame(
+            np.arange(27).reshape(3, 9),
+            columns=pd.MultiIndex.from_product(
+                [["a1", "a2", "a3"], ["b1", "b2", "b3"]]
+            ),
+        )
+        result = df1.loc(axis=1)["a1":"a2"]
+        expected = df1.iloc[:, :-3]
+        tm.assert_frame_equal(result, expected)
+        result = df1.loc(axis=1)["a1"]
+        expected = df1.iloc[:, :3]
+        expected.columns = ["b1", "b2", "b3"]
+        tm.assert_frame_equal(result, expected)
+
+        # test single level indexing on single index column data frame
+        df2 = pd.DataFrame(np.arange(9).reshape(3, 3), columns=["a", "b", "c"])
+        result = df2.loc(axis=1)["a"]
+        expected = pd.Series([0, 3, 6], dtype="int64", name="a")
+        tm.assert_series_equal(result, expected)
+
     def test_per_axis_per_level_setitem(self):
 
         # test index maker