Add restart grace period to multiball

mpf/config_spec.yaml

@@ -1141,6 +1141,7 @@ multiballs:
     shoot_again: single|template_ms|10s
     hurry_up_time: single|template_ms|0
     grace_period: single|template_ms|0
+    restart_grace_period: single|template_ms|0
     ball_locks: list|machine(ball_devices)|None
     enable_events: event_handler|event_handler:ms|None
     disable_events: event_handler|event_handler:ms|None

mpf/devices/multiball.py

@@ -237,15 +237,34 @@ def _ball_drain_count_balls(self, balls, **kwargs):
         desc: The multiball called (name) has lost a ball after ball save expired.
         '''
 
-        if not self.machine.game or self.machine.game.balls_in_play - balls < 1:
-            self.balls_added_live = 0
-            self.balls_live_target = 0
-            self.machine.events.remove_handler(self._ball_drain_count_balls)
-            self.machine.events.post("multiball_{}_ended".format(self.name))
-            '''event: multiball_(name)_ended
-            desc: The multiball called (name) has just ended.
-            '''
-            self.debug_log("Ball drained. MB ended.")
+        restart_grace_period_ms = self.config['restart_grace_period'].evaluate([])
+
+        if not self.machine.game or (self.machine.game.balls_in_play - balls < 1 and restart_grace_period_ms == 0):
+            self._multiball_end()
+        elif restart_grace_period_ms > 0:
+            if self.machine.game.balls_in_play == 0:
+                self.delay.run_now("restart_grace_period")
+            elif self.machine.game.balls_in_play - balls < 1:
+                self.machine.events.post(f"multiball_{self.name}_will_end", grace_period=restart_grace_period_ms)
+                '''event: multiball_(name)_will_end
+                desc: The multiball called (name) does not have multiple balls active, enabling grace period.
+                args:
+                    grace_period: restart duration grace period in ms
+                '''
+
+                self.delay.add(name="restart_grace_period",
+                               ms=restart_grace_period_ms,
+                               callback=self._multiball_end)
+
+    def _multiball_end(self):
+        self.balls_added_live = 0
+        self.balls_live_target = 0
+        self.machine.events.remove_handler(self._ball_drain_count_balls)
+        self.machine.events.post("multiball_{}_ended".format(self.name))
+        '''event: multiball_(name)_ended
+        desc: The multiball called (name) has just ended.
+        '''
+        self.debug_log("Ball drained. MB ended.")
 
     @event_handler(5)
     def event_stop(self, **kwargs):
@@ -284,6 +303,14 @@ def event_add_a_ball(self, **kwargs):
 
     def add_a_ball(self):
         """Add a ball if multiball has started."""
+        # remove if it exists as multiball now has multiple balls
+        if self.delay.check("restart_grace_period"):
+            self.delay.remove("restart_grace_period")
+            self.machine.events.post(f"multiball_{self.name}_will_resume")
+            '''event: multiball_(name)_will_resume
+            desc: Multiball (name) has been resumed during end of multiball grace period.
+            '''
+
         if self.balls_live_target > 0:
             self.debug_log("Adding a ball.")
             self.balls_live_target += 1

mpf/tests/machine_files/multiball/config/config.yaml

@@ -128,3 +128,12 @@ multiballs:
         add_a_ball_grace_period: 10s
         start_events: mb_add_a_ball_timers_start
         stop_events: mb_add_a_ball_timers_stop
+    mb_grace_period_restart:
+        ball_count: 2
+        add_a_ball_events: add_ball
+        add_a_ball_shoot_again: 20s
+        add_a_ball_hurry_up_time: 5s
+        add_a_ball_grace_period: 10s
+        restart_grace_period: 5s
+        start_events: mb_grace_period_restart_start
+        stop_events: mb_grace_period_restart_stop

mpf/tests/test_MultiBall.py

@@ -1228,3 +1228,58 @@ def testAddABallSaverDuringShootAgain(self):
         self.assertAvailableBallsOnPlayfield(3)
         self.assertEventCalled("ball_save_mb_add_a_ball_timers_timer_start", 1)
         self.assertEventNotCalled("ball_save_mb_add_a_ball_timers_add_a_ball_timer_start")
+
+    def testEndMultiballGracePeriodNoRestart(self):
+        self.fill_troughs()
+        self.start_game()
+        self.assertAvailableBallsOnPlayfield(1)
+        self.mock_event("multiball_mb_grace_period_restart_will_end")
+        self.mock_event("multiball_mb_grace_period_restart_ended")
+        self.mock_event("multiball_mb_grace_period_restart_will_resume")
+
+        # start mb, default ball save, 5s restart grace period
+        self.post_event("mb_grace_period_restart_start")
+        self.advance_time_and_run(10)
+        self.assertAvailableBallsOnPlayfield(2)
+
+        self.drain_one_ball()
+        self.advance_time_and_run(10)
+
+        self.assertEventCalled("multiball_mb_grace_period_restart_will_end", 1)
+        self.assertEventNotCalled("multiball_mb_grace_period_restart_will_resume")
+        self.assertEventCalled("multiball_mb_grace_period_restart_ended", 1)
+        self.assertAvailableBallsOnPlayfield(1)
+
+
+    def testEndMultiballGracePeriodAddABallRestart(self):
+        self.fill_troughs()
+        self.start_game()
+        self.assertAvailableBallsOnPlayfield(1)
+        self.mock_event("multiball_mb_grace_period_restart_will_end")
+        self.mock_event("multiball_mb_grace_period_restart_ended")
+        self.mock_event("multiball_mb_grace_period_restart_will_resume")
+        self.mock_event("ball_save_mb_grace_period_restart_timer_start")
+
+        # start mb, default ball save, 5s restart grace period
+        self.post_event("mb_grace_period_restart_start")
+        self.advance_time_and_run(10)
+        self.assertAvailableBallsOnPlayfield(2)
+
+        self.drain_one_ball()
+        self.advance_time_and_run(3)
+
+        # we should now be in the grace period (set at 5 seconds)
+        self.assertEventCalled("multiball_mb_grace_period_restart_will_end", 1)
+        self.assertEventNotCalled("multiball_mb_grace_period_restart_ended")
+
+        # trigger the add a ball during grace period
+        self.post_event("add_ball")
+        self.advance_time_and_run()
+        self.assertEventNotCalled("multiball_mb_grace_period_restart_ended")
+        self.assertEventCalled("multiball_mb_grace_period_restart_will_resume", 1)
+
+        # fast-forward time past the grace period and ensure multiball has resumed (i.e. with 2 balls)
+        self.advance_time_and_run(10)
+        self.assertAvailableBallsOnPlayfield(2)
+        self.assertEventCalled("ball_save_mb_grace_period_restart_timer_start")
+        self.assertEventNotCalled("multiball_mb_grace_period_restart_ended")