detect the "mandatory switch" menuvba-automation

This requires a slightly slower text_wait function. There is probably a
way to refactor that function in a way that doesn't cause cancer.

3 files changed, 31 insertions, 5 deletions

diff --git a/pokemontools/vba/battle.py b/pokemontools/vba/battle.py
index 824f27a..87cd7b1 100644
--- a/pokemontools/vba/battle.py
+++ b/pokemontools/vba/battle.py
@@ -67,8 +67,10 @@ class Battle(EmulatorController):
         #   1) current pokemon hp is 0
         #   2) game is polling for input
 
-        # TODO: detect the mandatory switch menu
-        return False
+        if "CANCEL Which ?" in self.emulator.get_text():
+            return True
+        else:
+            return False
 
     def skip_start_text(self, max_loops=20):
         """
diff --git a/pokemontools/vba/vba.py b/pokemontools/vba/vba.py
index 0dac63f..10513c6 100644
--- a/pokemontools/vba/vba.py
+++ b/pokemontools/vba/vba.py
@@ -431,12 +431,17 @@ class crystal(object):
                 # set CurSFX
                 self.vba.write_memory_at(0xc2bf, 0)
 
-                self.vba.press("a", hold=10, after=1)
+                self.vba.press("a", hold=10, after=50)
 
                 # check if CurSFX is SFX_READ_TEXT_2
                 if self.vba.read_memory_at(0xc2bf) == 0x8:
-                    print "cursfx is set to SFX_READ_TEXT_2, looping.."
-                    return self.text_wait(step_size=step_size, max_wait=max_wait, debug=debug, callback=callback, sfx_limit=sfx_limit)
+                    if "CANCEL Which" in self.get_text():
+                        print "probably the 'switch pokemon' menu"
+                        return
+                    else:
+                        print "cursfx is set to SFX_READ_TEXT_2, looping.."
+                        print self.get_text()
+                        return self.text_wait(step_size=step_size, max_wait=max_wait, debug=debug, callback=callback, sfx_limit=sfx_limit)
                 else:
                     if sfx_limit > 0:
                         sfx_limit = sfx_limit - 1
diff --git a/tests/test_vba_battle.py b/tests/test_vba_battle.py
index e86b53c..61c0297 100644
--- a/tests/test_vba_battle.py
+++ b/tests/test_vba_battle.py
@@ -56,6 +56,25 @@ class BattleTests(unittest.TestCase):
         # should not be asking for a switch so soon in the battle
         self.assertFalse(self.battle.is_mandatory_switch())
 
+    def test_is_mandatory_switch(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 partymon1 hp to very low
+        self.vba.write_memory_at(0xc63c, 0)
+        self.vba.write_memory_at(0xc63d, 1)
+
+        # let the enemy attack and kill the pokemon
+        self.battle.skip_until_input_required()
+
+        self.assertTrue(self.battle.is_mandatory_switch())
+
     def test_attack_loop(self):
         self.battle.skip_start_text()
         self.battle.skip_until_input_required()