Recover from occupied bed race

mosfet/jobs/sleep_with_bed.py

@@ -32,6 +32,7 @@ class SleepWithBedStates:
         self.waiting_for_time_update = False
 
         current_time = self.g.get('time')
+        self.sleep_attempt_started_at = current_time
         if current_time is None or current_time < 12000:
             print('Aborting sleep, not night')
             self.state = self.cleanup
@@ -268,7 +269,22 @@ class SleepWithBedStates:
                 print('Waiting for world time update before using bed')
                 self.waiting_for_time_update = True
             return
+        if (self.area is not None and
+                w.block_at(*self.area) not in blocks.BED_IDS):
+            print('Bed disappeared before use:', self.area)
+            self.state = self.choose_bed
+            return
+        if self.area is not None and w.check_bed_occupied(self.area):
+            print('Bed became occupied before use:', self.area)
+            self.state = self.choose_bed
+            return
         if current_time < 12550:
+            if (self.sleep_attempt_started_at is not None and
+                    self.sleep_attempt_started_at >= 12000 and
+                    current_time < 12000):
+                print('Night ended before bed use; resuming the job')
+                self.state = self.cleanup
+                return
             if not getattr(self, 'waiting_for_night', False):
                 print('Bed reached; waiting until sleep time (time={})'.format(
                     current_time))
@@ -277,16 +293,6 @@ class SleepWithBedStates:
         self.waiting_for_night = False
         self.waiting_for_time_update = False
 
-        if w.block_at(*self.area) not in blocks.BED_IDS:
-            print('Bed disappeared before use:', self.area)
-            self.my_bed = False
-            self.state = self.cleanup
-            return
-        if w.check_bed_occupied(self.area):
-            print('Bed occupied, aborting.')
-            self.state = self.cleanup
-            return
-
         print('Using bed')
         # Beds are only 9/16 block tall. Aim inside their collision shape;
         # the near-top cursor used for placing a new block misses them.
@@ -404,6 +410,7 @@ class SleepWithBedStates:
         self.force_new_bed = False
         self.sleep_confirmed = False
         self.sleep_deadline = 0
+        self.sleep_attempt_started_at = None
 
     def run(self):
         self.state()

tests/test_sleep_with_bed.py

@@ -366,6 +366,42 @@ class TestSleepWithBed(unittest.TestCase):
         self.assertEqual(state.state, state.use_bed)
         self.assertTrue(state.waiting_for_night)
 
+    def test_sleep_attempt_aborts_if_night_was_skipped_while_approaching(self):
+        bed_id = next(iter(blocks.BED_IDS))
+        world = FakeWorld(block=bed_id)
+        state = SleepWithBedStates(Munch(
+            time=13000, dimension='overworld', world=world,
+            game=FakeGame(), path=[], look_at=None))
+        state.init()
+        state.area = (0, 64, 0)
+        state.state = state.use_bed
+        state.g.time = 25
+
+        state.use_bed()
+
+        self.assertEqual(state.state, state.cleanup)
+
+    def test_occupied_bed_is_abandoned_before_waiting_for_night(self):
+        bed_id = next(iter(blocks.BED_IDS))
+
+        class OccupiedWorld(FakeWorld):
+            @staticmethod
+            def check_bed_occupied(_position):
+                return True
+
+        state = SleepWithBedStates(Munch(
+            time=13000, dimension='overworld',
+            world=OccupiedWorld(block=bed_id), game=FakeGame(),
+            path=[], look_at=None))
+        state.init()
+        state.area = (0, 64, 0)
+        state.beds = [(2, 64, 0)]
+        state.state = state.use_bed
+
+        state.use_bed()
+
+        self.assertEqual(state.state, state.choose_bed)
+
     def test_daytime_sleep_aborts_cleanly(self):
         state = SleepWithBedStates(Munch(time=1000, dimension='overworld'))
         state.init()