diff --git a/pandas/tests/frame/methods/test_at_time.py b/pandas/tests/frame/methods/test_at_time.py
index 2d05176d20f5f..e15a5f80e5008 100644
--- a/pandas/tests/frame/methods/test_at_time.py
+++ b/pandas/tests/frame/methods/test_at_time.py
@@ -101,12 +101,8 @@ def test_at_time_axis(self, axis):
         ts.index, ts.columns = rng, rng
 
         indices = rng[(rng.hour == 9) & (rng.minute == 30) & (rng.second == 0)]
-
-        if axis in ["index", 0]:
-            expected = ts.loc[indices, :]
-        elif axis in ["columns", 1]:
-            expected = ts.loc[:, indices]
-
+        axis = self.data._get_axis_number(axis)
+        expected = ts.loc[indices, :] if axis == 0 else ts.loc[:, indices]
         result = ts.at_time("9:30", axis=axis)
 
         # Without clearing freq, result has freq 1440T and expected 5T