From 2af6d65893dbd37c1f370f3d71ef1c7a73a3e840 Mon Sep 17 00:00:00 2001
From: fakuzatsu <118256341+fakuzatsu@users.noreply.github.com>
Date: Thu, 21 Dec 2023 13:01:13 +0000
Subject: [PATCH] added ability to set flags in tests (#3786)

---
 include/test/battle.h     | 13 +++++++++++++
 test/test_runner_battle.c | 20 ++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/include/test/battle.h b/include/test/battle.h
index cb62e5b56f..32908df9c5 100644
--- a/include/test/battle.h
+++ b/include/test/battle.h
@@ -280,6 +280,13 @@
  *     GIVEN {
  *         RNGSeed(0xC0DEIDEA);
  *
+ * FLAG_SET(flagId)
+ * Sets the specified flag. Can currently only set one flag at a time.
+ * Cleared between perameters and at the end of the test.
+ * Example:
+ *     GIVEN {
+ *         FLAG_SET(FLAG_SYS_EXAMPLE_FLAG);
+ *
  * PLAYER(species) and OPPONENT(species)
  * Adds the species to the player's or opponent's party respectively.
  * The Pokémon can be further customized with the following functions:
@@ -489,6 +496,7 @@
 #include "constants/battle_ai.h"
 #include "constants/battle_anim.h"
 #include "constants/battle_move_effects.h"
+#include "constants/flags.h"
 #include "constants/hold_effects.h"
 #include "constants/items.h"
 #include "constants/moves.h"
@@ -681,6 +689,7 @@ struct BattleTestData
     struct ExpectedAiScore expectedAiScores[MAX_BATTLERS_COUNT][MAX_TURNS][MAX_AI_SCORE_COMPARISION_PER_TURN]; // Max 4 comparisions per turn
     struct AILogLine aiLogLines[MAX_BATTLERS_COUNT][MAX_MON_MOVES][MAX_AI_LOG_LINES];
     u8 aiLogPrintedForMove[MAX_BATTLERS_COUNT]; // Marks ai score log as printed for move, so the same log isn't displayed multiple times.
+    u16 flagId;
 };
 
 struct BattleTestRunnerState
@@ -818,6 +827,8 @@ struct moveWithPP {
 #define AI_FLAGS(flags) AIFlags_(__LINE__, flags)
 #define AI_LOG AILogScores(__LINE__)
 
+#define FLAG_SET(flagId) SetFlagForTest(__LINE__, flagId)
+
 #define PLAYER(species) for (OpenPokemon(__LINE__, B_SIDE_PLAYER, species); gBattleTestRunnerState->data.currentMon; ClosePokemon(__LINE__))
 #define OPPONENT(species) for (OpenPokemon(__LINE__, B_SIDE_OPPONENT, species); gBattleTestRunnerState->data.currentMon; ClosePokemon(__LINE__))
 
@@ -839,6 +850,8 @@ struct moveWithPP {
 #define Status1(status1) Status1_(__LINE__, status1)
 #define OTName(otName) do {static const u8 otName_[] = _(otName); OTName_(__LINE__, otName_);} while (0)
 
+void SetFlagForTest(u32 sourceLine, u16 flagId);
+void ClearFlagAfterTest(void);
 void OpenPokemon(u32 sourceLine, u32 side, u32 species);
 void ClosePokemon(u32 sourceLine);
 
diff --git a/test/test_runner_battle.c b/test/test_runner_battle.c
index ac689676b0..d34e1fc96c 100644
--- a/test/test_runner_battle.c
+++ b/test/test_runner_battle.c
@@ -4,6 +4,7 @@
 #include "battle_anim.h"
 #include "battle_controllers.h"
 #include "characters.h"
+#include "event_data.h"
 #include "fieldmap.h"
 #include "item_menu.h"
 #include "main.h"
@@ -1339,6 +1340,7 @@ static void CB2_BattleTest_NextParameter(void)
     if (++STATE->runParameter >= STATE->parameters)
     {
         SetMainCallback2(CB2_TestRunner);
+        ClearFlagAfterTest();
     }
     else
     {
@@ -1349,6 +1351,7 @@ static void CB2_BattleTest_NextParameter(void)
 
 static void CB2_BattleTest_NextTrial(void)
 {
+    ClearFlagAfterTest();
     TearDownBattle();
 
     SetMainCallback2(CB2_BattleTest_NextParameter);
@@ -1390,6 +1393,7 @@ static void BattleTest_TearDown(void *data)
 {
     // Free resources that aren't cleaned up when the battle was
     // aborted unexpectedly.
+    ClearFlagAfterTest();
     if (STATE->tearDownBattle)
         TearDownBattle();
 }
@@ -1478,6 +1482,22 @@ const struct TestRunner gBattleTestRunner =
     .handleExitWithResult = BattleTest_HandleExitWithResult,
 };
 
+void SetFlagForTest(u32 sourceLine, u16 flagId)
+{
+    INVALID_IF(DATA.flagId != 0, "FLAG can only be set once per test");
+    DATA.flagId = flagId;
+    FlagSet(flagId);
+}
+
+void ClearFlagAfterTest(void)
+{
+    if (DATA.flagId != 0) 
+    {
+        FlagClear(DATA.flagId);
+        DATA.flagId = 0;
+    }
+}
+
 void OpenPokemon(u32 sourceLine, u32 side, u32 species)
 {
     s32 i, data;