# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2026) # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from __future__ import annotations import asyncio import os import time import types from typing import TYPE_CHECKING, Any from urllib import parse from streamlit import runtime from streamlit.runtime.forward_msg_queue import ForwardMsgQueue from streamlit.runtime.fragment import MemoryFragmentStorage from streamlit.runtime.memory_uploaded_file_manager import MemoryUploadedFileManager from streamlit.runtime.scriptrunner import RerunData, ScriptRunner, ScriptRunnerEvent from streamlit.runtime.scriptrunner.script_cache import ScriptCache from streamlit.testing.v1.element_tree import ElementTree, parse_tree_from_messages if TYPE_CHECKING: from streamlit.proto.ForwardMsg_pb2 import ForwardMsg from streamlit.proto.WidgetStates_pb2 import WidgetStates from streamlit.runtime.pages_manager import PagesManager from streamlit.runtime.scriptrunner_utils.script_run_context import ScriptRunContext from streamlit.runtime.state.safe_session_state import SafeSessionState from streamlit.runtime.uploaded_file_manager import UploadedFileRec class LocalScriptRunner(ScriptRunner): """AppTest adapter that captures events and owns each run's event loop.""" def __init__( self, script_path: str, session_state: SafeSessionState, pages_manager: PagesManager, args: Any = None, kwargs: Any = None, fragment_storage: MemoryFragmentStorage | None = None, ) -> None: """Initializes the ScriptRunner for the given script_path.""" if not os.path.isfile(script_path): raise FileNotFoundError(f"File not found at {script_path}") self.forward_msg_queue = ForwardMsgQueue() self.script_path = script_path self.session_state = session_state self.args = args if args is not None else () self.kwargs = kwargs if kwargs is not None else {} self._owned_event_loop = asyncio.new_event_loop() try: super().__init__( session_id="test session id", main_script_path=script_path, session_state=self.session_state._state, uploaded_file_mgr=MemoryUploadedFileManager("/mock/upload"), script_cache=ScriptCache(), initial_rerun_data=RerunData(), user_info={"email": "test@example.com"}, fragment_storage=( fragment_storage if fragment_storage is not None else MemoryFragmentStorage() ), pages_manager=pages_manager, event_loop=self._owned_event_loop, ) except Exception: self._owned_event_loop.close() raise # Accumulates all ScriptRunnerEvents emitted by us. self.events: list[ScriptRunnerEvent] = [] self.event_data: list[Any] = [] def record_event( sender: ScriptRunner | None, event: ScriptRunnerEvent, **kwargs: Any ) -> None: # Assert that we're not getting unexpected `sender` params # from ScriptRunner.on_event if sender is not None and sender != self: raise RuntimeError("Unexpected ScriptRunnerEvent sender!") self.events.append(event) self.event_data.append(kwargs) if event == ScriptRunnerEvent.SCRIPT_STARTED: # Drop stale deltas from the previous run so the element tree # does not accumulate elements across (internal) reruns. This # mirrors AppSession._clear_queue: fragment_ids_this_run is # forwarded so that a fragment-scoped rerun preserves deltas # belonging to elements outside the running fragment(s). self.forward_msg_queue.clear( retain_lifecycle_msgs=True, fragment_ids_this_run=kwargs.get("fragment_ids_this_run"), ) # Send ENQUEUE_FORWARD_MSGs to our queue if event == ScriptRunnerEvent.ENQUEUE_FORWARD_MSG: forward_msg = kwargs["forward_msg"] self.forward_msg_queue.enqueue(forward_msg) self.on_event.connect(record_event, weak=False) def register_file(self, file_rec: UploadedFileRec) -> None: """Register an uploaded file with the file manager. This is a typed helper method to avoid accessing private attributes directly from AppTest. Parameters ---------- file_rec The uploaded file record to register. """ # Cast to MemoryUploadedFileManager since that's what LocalScriptRunner uses. # The Protocol UploadedFileManager doesn't include add_file. from typing import cast file_mgr = cast("MemoryUploadedFileManager", self._uploaded_file_mgr) file_mgr.add_file(self._session_id, file_rec) def join(self) -> None: """Wait for the script thread to finish, if it is running.""" if self._script_thread is not None: self._script_thread.join() self._close_owned_event_loop() def _run_script_thread(self) -> None: try: super()._run_script_thread() finally: self._close_owned_event_loop() def _close_owned_event_loop(self) -> None: """Close the test-owned loop after the script thread releases it.""" if not self._owned_event_loop.is_closed(): self._owned_event_loop.close() def forward_msgs(self) -> list[ForwardMsg]: """Return all messages in our ForwardMsgQueue.""" return self.forward_msg_queue._queue def run( self, widget_state: WidgetStates | None = None, query_params: dict[str, Any] | None = None, timeout: float = 3, page_hash: str = "", ) -> ElementTree: """Run the script, and parse the output messages for querying and interaction. Timeout is in seconds. """ # TODO: save the query strings from the script run query_string = "" if query_params: query_string = parse.urlencode(query_params, doseq=True) rerun_data = RerunData( widget_states=widget_state, query_string=query_string, page_script_hash=page_hash, ) self.request_rerun(rerun_data) try: if not self._script_thread: self.start() require_widgets_deltas(self, timeout) finally: # SHUTDOWN is emitted before the script thread returns. Joining # ensures AppTest never observes completion before loop cleanup. self.join() return parse_tree_from_messages(self.forward_msgs()) def script_stopped(self) -> bool: return any(e == ScriptRunnerEvent.SHUTDOWN for e in self.events) def _on_script_finished( self, ctx: ScriptRunContext, event: ScriptRunnerEvent, premature_stop: bool ) -> None: # Do not drop unseen widgets when the script stopped for a rerun # (later widgets never registered). Keep this gate in sync with # ScriptRunner._on_script_finished. if not premature_stop: self._session_state.on_script_finished( ctx.shared.widget_ids_this_run.snapshot(), remove_stale_widgets=( ctx.has_script_started and event != ScriptRunnerEvent.SCRIPT_STOPPED_FOR_RERUN ), ) # Signal that the script has finished. (We use SCRIPT_STOPPED_WITH_SUCCESS # even if we were stopped with an exception.) self.on_event.send(self, event=event) # Skip orphan cleanup when: # - The body never ran (nothing was re-registered, so everything # would look orphaned) # - No Runtime exists (get_instance() would raise) if ctx.has_script_started and runtime.exists(): runtime.get_instance().media_file_mgr.remove_orphaned_files() def _new_module(self, name: str) -> types.ModuleType: module = types.ModuleType(name) module.__dict__["__args"] = self.args module.__dict__["__kwargs"] = self.kwargs return module def require_widgets_deltas(runner: LocalScriptRunner, timeout: float = 3) -> None: """Wait for the given ScriptRunner to emit a completion event. If the timeout is reached, the runner will be shutdown and an error will be thrown. """ t0 = time.time() while time.time() - t0 < timeout: time.sleep(0.001) if runner.script_stopped(): return # If we get here, the runner hasn't yet completed before our # timeout. Create an error string for debugging. err_string = f"AppTest script run timed out after {timeout}(s)" # Shutdown the runner before throwing an error, so that the script # doesn't hang forever. runner.request_stop() runner.join() raise RuntimeError(err_string)