From 7e56e94299e15c0af87aabe269da4ef55ee383f7 Mon Sep 17 00:00:00 2001
From: Vishal <vishal.u338@gmail.com>
Date: Sat, 24 Sep 2022 21:52:24 +0530
Subject: [PATCH] End Test Text instead of End Game for Stress Tests

---
 assets/src/ba_data/python/ba/_benchmark.py    |  5 ++--
 .../python/bastd/ui/account/settings.py       |  2 +-
 .../src/ba_data/python/bastd/ui/mainmenu.py   | 23 +++++++++++++++++--
 3 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/assets/src/ba_data/python/ba/_benchmark.py b/assets/src/ba_data/python/ba/_benchmark.py
index 119296c9..53ea3f5d 100644
--- a/assets/src/ba_data/python/ba/_benchmark.py
+++ b/assets/src/ba_data/python/ba/_benchmark.py
@@ -60,7 +60,7 @@ def run_stress_test(playlist_type: str = 'Random',
     from ba._generated.enums import TimeType
     _ba.screenmessage(
         'Beginning stress test.. use '
-        "'End Game' to stop testing.",
+        "'End Test' to stop testing.",
         color=(1, 1, 0))
     with _ba.Context('ui'):
         start_stress_test({
@@ -138,7 +138,8 @@ def _reset_stress_test(args: dict[str, Any]) -> None:
 
 def run_gpu_benchmark() -> None:
     """Kick off a benchmark to test gpu speeds."""
-    _ba.screenmessage('FIXME: Not wired up yet.', color=(1, 0, 0))
+    # FIXME: Not wired up yet.
+    _ba.screenmessage('Not wired up yet.', color=(1, 0, 0))
 
 
 def run_media_reload_benchmark() -> None:
diff --git a/assets/src/ba_data/python/bastd/ui/account/settings.py b/assets/src/ba_data/python/bastd/ui/account/settings.py
index a75eb0f1..1fdc3b7f 100644
--- a/assets/src/ba_data/python/bastd/ui/account/settings.py
+++ b/assets/src/ba_data/python/bastd/ui/account/settings.py
@@ -552,7 +552,7 @@ class AccountSettingsWindow(ba.Window):
                 position=((self._sub_width - button_width) * 0.5, v + 30),
                 autoselect=True,
                 size=(button_width, 60),
-                label=ba.Lstr(resource=self._r + '.manageAccount'),
+                label=ba.Lstr(resource=self._r + '.manageAccountText'),
                 color=(0.55, 0.5, 0.6),
                 icon=ba.gettexture('settingsIcon'),
                 textcolor=(0.75, 0.7, 0.8),
diff --git a/assets/src/ba_data/python/bastd/ui/mainmenu.py b/assets/src/ba_data/python/bastd/ui/mainmenu.py
index 6e8ad23e..39f38371 100644
--- a/assets/src/ba_data/python/bastd/ui/mainmenu.py
+++ b/assets/src/ba_data/python/bastd/ui/mainmenu.py
@@ -250,8 +250,12 @@ class MainMenuWindow(ba.Window):
                     scale=scale,
                     size=(self._button_width, self._button_height),
                     autoselect=self._use_autoselect,
-                    label=ba.Lstr(resource=self._r + '.endGameText'),
-                    on_activate_call=self._confirm_end_game)
+                    label=ba.Lstr(resource=self._r +
+                                  ('.endTestText' if self._is_benchmark()
+                                   else '.endGameText')),
+                    on_activate_call=(self._confirm_end_test
+                                      if self._is_benchmark()
+                                      else self._confirm_end_game))
             # Assume we're in a client-session.
             else:
                 ba.buttonwidget(
@@ -863,6 +867,11 @@ class MainMenuWindow(ba.Window):
             StoreBrowserWindow(
                 origin_widget=self._store_button).get_root_widget())
 
+    def _is_benchmark(self) -> bool:
+        session = ba.internal.get_foreground_host_session()
+        return (getattr(session, 'benchmark_type', None) == 'cpu'
+            or ba.app.stress_test_reset_timer is not None)
+
     def _confirm_end_game(self) -> None:
         # pylint: disable=cyclic-import
         from bastd.ui.confirm import ConfirmWindow
@@ -874,6 +883,16 @@ class MainMenuWindow(ba.Window):
                       self._end_game,
                       cancel_is_selected=True)
 
+    def _confirm_end_test(self) -> None:
+        # pylint: disable=cyclic-import
+        from bastd.ui.confirm import ConfirmWindow
+
+        # Select cancel by default; this occasionally gets called by accident
+        # in a fit of button mashing and this will help reduce damage.
+        ConfirmWindow(ba.Lstr(resource=self._r + '.exitToMenuText'),
+                      self._end_game,
+                      cancel_is_selected=True)
+
     def _confirm_end_replay(self) -> None:
         # pylint: disable=cyclic-import
         from bastd.ui.confirm import ConfirmWindow