Merge pull request #52 from kanzure/battle-handling

Battles

1 files changed, 26 insertions, 2 deletions

diff --git a/tests/test_vba_battle.py b/tests/test_vba_battle.py
index 61c0297..c6debc3 100644
--- a/tests/test_vba_battle.py
+++ b/tests/test_vba_battle.py
@@ -67,8 +67,7 @@ class BattleTests(unittest.TestCase):
         self.vba.press(["a"], after=20)
 
         # set partymon1 hp to very low
-        self.vba.write_memory_at(0xc63c, 0)
-        self.vba.write_memory_at(0xc63d, 1)
+        self.cry.set_battle_mon_hp(1)
 
         # let the enemy attack and kill the pokemon
         self.battle.skip_until_input_required()
@@ -89,5 +88,30 @@ class BattleTests(unittest.TestCase):
 
         self.assertTrue(self.battle.is_player_turn())
 
+    def test_is_battle_switch_prompt(self):
+        self.battle.skip_start_text()
+        self.battle.skip_until_input_required()
+
+        # press "FIGHT"
+        self.vba.press(["a"], after=20)
+
+        # press the first move ("SCRATCH")
+        self.vba.press(["a"], after=20)
+
+        # set enemy hp to very low
+        self.cry.lower_enemy_hp()
+
+        # attack the enemy and kill it
+        self.battle.skip_until_input_required()
+
+        # yes/no menu is present, should be detected
+        self.assertTrue(self.battle.is_trainer_switch_prompt())
+
+        # and input should be required
+        self.assertTrue(self.battle.is_input_required())
+
+        # but it's not mandatory
+        self.assertFalse(self.battle.is_mandatory_switch())
+
 if __name__ == "__main__":
     unittest.main()