diff --git a/test/test_esm.py b/test/test_esm.py
index 880daa2..6ee7513 100644
--- a/test/test_esm.py
+++ b/test/test_esm.py
@@ -25,9 +25,10 @@ class EsmFeatures(unittest.TestCase):
 
     def test_classify_sessions_by_completion_time(self):
         self.esm_classified = classify_sessions_by_completion_time(self.esm_processed)
-        session_of_interest = self.esm_classified.query("(device_id == '049df3f8-8541-4cf5-af2b-83f6b3f0cf4b') & (esm_session == 1)")
-        print(session_of_interest)
-        self.assertEqual(session_of_interest["time"], "daytime")
+        session_of_interest = self.esm_processed.query("(device_id == '049df3f8-8541-4cf5-af2b-83f6b3f0cf4b') & (esm_session == 1)").sort_values("_id").reset_index()
+        session_of_interest_reclassified = self.esm_classified.query("(device_id == '049df3f8-8541-4cf5-af2b-83f6b3f0cf4b') & (esm_session == 1)").reset_index()
+        self.assertEqual(session_of_interest.loc[0, "time"], "morning")
+        self.assertEqual(session_of_interest_reclassified.loc[0, "time"], "daytime")
         # Check that the first (morning) session is reclassified as a daytime session.
 
     def test_clean_up_esm(self):