Fix infinite loop when rerun and triggered widgets are used togethe…

…r in AppTest (#8264)

## Describe your changes
Reverts a previous fix that made trigger values correct in AppTest but
introduced an infinite recursion when `st.rerun()` is called
conditionally on the value of a trigger widget. Instead, making trigger
widgets return correct values in AppTest is handled by storing their
values at runtime using the new `save_for_app_testing` functionality
introduced by #8189

## GitHub Issue Link (if applicable)
#7768

lib/streamlit/elements/widgets/button.py

@@ -39,7 +39,7 @@
     WidgetKwargs,
     register_widget,
 )
-from streamlit.runtime.state.common import compute_widget_id
+from streamlit.runtime.state.common import compute_widget_id, save_for_app_testing
 from streamlit.string_util import validate_emoji
 from streamlit.type_util import Key, to_key
 
@@ -762,6 +762,8 @@ def _button(
             ctx=ctx,
         )
 
+        if ctx:
+            save_for_app_testing(ctx, id, button_state.value)
         self.dg._enqueue("button", button_proto)
 
         return button_state.value

lib/streamlit/elements/widgets/chat.py

@@ -35,7 +35,7 @@
     WidgetKwargs,
     register_widget,
 )
-from streamlit.runtime.state.common import compute_widget_id
+from streamlit.runtime.state.common import compute_widget_id, save_for_app_testing
 from streamlit.string_util import is_emoji
 from streamlit.type_util import Key, to_key
 
@@ -352,6 +352,8 @@ def chat_input(
             chat_input_proto.value = widget_state.value
             chat_input_proto.set_value = True
 
+        if ctx:
+            save_for_app_testing(ctx, id, widget_state.value)
         if position == "bottom":
             # We import it here to avoid circular imports.
             from streamlit import _bottom

lib/streamlit/testing/v1/element_tree.py

@@ -334,7 +334,7 @@ def value(self) -> bool:
         else:
             state = self.root.session_state
             assert state
-            return cast(bool, state[self.id])
+            return cast(bool, state[TESTING_KEY][self.id])
 
     def set_value(self, v: bool) -> Button:
         """Set the value of the button."""
@@ -379,7 +379,7 @@ def value(self) -> str | None:
         else:
             state = self.root.session_state
             assert state
-            return state[self.id]  # type: ignore
+            return state[TESTING_KEY][self.id]  # type: ignore
 
 
 @dataclass(repr=False)

lib/streamlit/testing/v1/local_script_runner.py

@@ -127,10 +127,8 @@ def script_stopped(self) -> bool:
     def _on_script_finished(
         self, ctx: ScriptRunContext, event: ScriptRunnerEvent, premature_stop: bool
     ) -> None:
-        # Only call `_remove_stale_widgets`, so that the state of triggers is still
-        # visible in the element tree.
         if not premature_stop:
-            self._session_state._state._remove_stale_widgets(ctx.widget_ids_this_run)
+            self._session_state.on_script_finished(ctx.widget_ids_this_run)
 
         # Signal that the script has finished. (We use SCRIPT_STOPPED_WITH_SUCCESS
         # even if we were stopped with an exception.)

lib/tests/streamlit/testing/app_test_test.py

@@ -217,3 +217,20 @@ def script(foo, baz):
 
     at = AppTest.from_function(script, args=("bar",), kwargs={"baz": "baz"}).run()
     assert at.text.values == ["bar", "baz"]
+
+
+def test_trigger_recursion():
+    # Regression test for #7768
+    def code():
+        import time
+
+        import streamlit as st
+
+        if st.button(label="Submit"):
+            print("CLICKED!")
+            time.sleep(1)
+            st.rerun()
+
+    at = AppTest.from_function(code).run()
+    # The script run should finish instead of recurring and timing out
+    at.button[0].click().run()