From 953f2292e2b28a031b02c5171be7b901515a0b89 Mon Sep 17 00:00:00 2001
From: Pawkkie <61265402+Pawkkie@users.noreply.github.com>
Date: Sat, 30 Nov 2024 17:39:27 -0500
Subject: [PATCH] Fixes Switch in flag not restoring mons properly with test
 (#5746)

---
 src/battle_ai_util.c          |  1 +
 test/battle/ai/ai_switching.c | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/src/battle_ai_util.c b/src/battle_ai_util.c
index 0277d7fd07..c033061dc9 100644
--- a/src/battle_ai_util.c
+++ b/src/battle_ai_util.c
@@ -3507,6 +3507,7 @@ u32 AI_WhoStrikesFirstPartyMon(u32 battlerAtk, u32 battlerDef, struct BattlePoke
     SetBattlerAiData(battlerAtk, AI_DATA);
     u32 aiMonFaster = AI_IsFaster(battlerAtk, battlerDef, moveConsidered);
     FreeRestoreBattleMons(savedBattleMons);
+    SetBattlerAiData(battlerAtk, AI_DATA);
 
     return aiMonFaster;
 }
diff --git a/test/battle/ai/ai_switching.c b/test/battle/ai/ai_switching.c
index 5610689f49..7f2368261d 100644
--- a/test/battle/ai/ai_switching.c
+++ b/test/battle/ai/ai_switching.c
@@ -895,3 +895,17 @@ AI_SINGLE_BATTLE_TEST("Switch AI: AI will switch into mon with good type matchup
         TURN { MOVE(player, MOVE_TACKLE); EXPECT_SWITCH(opponent, 1); }
     }
 }
+
+AI_SINGLE_BATTLE_TEST("AI_FLAG_SMART_MON_CHOICES: AI correctly handles abilities when scoring moves")
+{
+    GIVEN {
+        ASSUME(B_PRANKSTER_DARK_TYPES >= GEN_7);
+        ASSUME(gSpeciesInfo[SPECIES_GRENINJA].types[1] == TYPE_DARK);
+        AI_FLAGS(AI_FLAG_CHECK_BAD_MOVE | AI_FLAG_TRY_TO_FAINT | AI_FLAG_CHECK_VIABILITY | AI_FLAG_OMNISCIENT | AI_FLAG_SMART_MON_CHOICES);
+        PLAYER(SPECIES_GRENINJA) { Moves(MOVE_WATER_GUN); }
+        OPPONENT(SPECIES_WHIMSICOTT) { Ability(ABILITY_PRANKSTER); Moves(MOVE_LEECH_SEED, MOVE_STUN_SPORE, MOVE_ABSORB); }
+        OPPONENT(SPECIES_WHIMSICOTT) { Ability(ABILITY_INFILTRATOR); }
+    } WHEN {
+        TURN { MOVE(player, MOVE_WATER_GUN); EXPECT_MOVE(opponent, MOVE_ABSORB); }
+    }
+}