def test_enqueue_new_session_message(self, _1, _2, patched_config):
def get_option(name):
if name == "server.runOnSave":
# Just to avoid starting the watcher for no reason.
return False
return config.get_option(name)
patched_config.get_option.side_effect = get_option
patched_config.get_options_for_section.side_effect = (
_mock_get_options_for_section()
)
# Create a AppSession with some mocked bits
rs = AppSession(
self.io_loop,
SessionData("mock_report.py", ""),
UploadedFileManager(),
lambda: None,
MagicMock(),
)
orig_ctx = get_script_run_ctx()
ctx = ScriptRunContext(
"TestSessionID", rs._session_data.enqueue, "", None, None
)
add_script_run_ctx(ctx=ctx)
rs._on_scriptrunner_event(ScriptRunnerEvent.SCRIPT_STARTED)
sent_messages = rs._session_data._browser_queue._queue
self.assertEqual(len(sent_messages), 2) # NewApp and SessionState messages
# Note that we're purposefully not very thoroughly testing new_session
# fields below to avoid getting to the point where we're just
# duplicating code in tests.
new_session_msg = sent_messages[0].new_session
self.assertEqual("mock_scriptrun_id", new_session_msg.script_run_id)
self.assertEqual(new_session_msg.HasField("config"), True)
self.assertEqual(
new_session_msg.config.allow_run_on_save,
config.get_option("server.allowRunOnSave"),
)
self.assertEqual(new_session_msg.HasField("custom_theme"), True)
self.assertEqual(new_session_msg.custom_theme.text_color, "black")
init_msg = new_session_msg.initialize
self.assertEqual(init_msg.HasField("user_info"), True)
add_script_run_ctx(ctx=orig_ctx)
def _run_script_thread(self) -> None:
"""The entry point for the script thread.
Processes the ScriptRequestQueue, which will at least contain the RERUN
request that will trigger the first script-run.
When the ScriptRequestQueue is empty, or when a SHUTDOWN request is
dequeued, this function will exit and its thread will terminate.
"""
assert self._is_in_script_thread()
LOGGER.debug("Beginning script thread")
# Create and attach the thread's ScriptRunContext
ctx = ScriptRunContext(
session_id=self._session_id,
enqueue=self._enqueue,
query_string=self._client_state.query_string,
session_state=self._session_state,
uploaded_file_mgr=self._uploaded_file_mgr,
)
add_script_run_ctx(threading.current_thread(), ctx)
while not self._shutdown_requested and self._request_queue.has_request:
request, data = self._request_queue.dequeue()
if request == ScriptRequest.STOP:
LOGGER.debug("Ignoring STOP request while not running")
elif request == ScriptRequest.SHUTDOWN:
LOGGER.debug("Shutting down")
self._shutdown_requested = True
elif request == ScriptRequest.RERUN:
self._run_script(data)
else:
raise RuntimeError("Unrecognized ScriptRequest: %s" % request)
# Send a SHUTDOWN event before exiting. This includes the widget values
# as they existed after our last successful script run, which the
# AppSession will pass on to the next ScriptRunner that gets
# created.
client_state = ClientState()
client_state.query_string = ctx.query_string
widget_states = self._session_state.as_widget_states()
client_state.widget_states.widgets.extend(widget_states)
self.on_event.send(self,
event=ScriptRunnerEvent.SHUTDOWN,
client_state=client_state)
def setUp(self, override_root=True):
self.forward_msg_queue = ForwardMsgQueue()
self.override_root = override_root
self.orig_report_ctx = None
self.new_script_run_ctx = ScriptRunContext(
session_id="test session id",
enqueue=self.forward_msg_queue.enqueue,
query_string="",
session_state=SessionState(),
uploaded_file_mgr=UploadedFileManager(),
)
if self.override_root:
self.orig_report_ctx = get_script_run_ctx()
add_script_run_ctx(threading.current_thread(),
self.new_script_run_ctx)
self.app_session = FakeAppSession()
def test_disallow_set_page_config_twice(self):
"""st.set_page_config cannot be called twice"""
fake_enqueue = lambda msg: None
ctx = ScriptRunContext(
"TestSessionID",
fake_enqueue,
"",
SessionState(),
UploadedFileManager(),
)
ctx.on_script_start()
msg = ForwardMsg()
msg.page_config_changed.title = "foo"
ctx.enqueue(msg)
with self.assertRaises(StreamlitAPIException):
same_msg = ForwardMsg()
same_msg.page_config_changed.title = "bar"
ctx.enqueue(same_msg)
def test_set_page_config_first(self):
"""st.set_page_config must be called before other st commands
when the script has been marked as started"""
fake_enqueue = lambda msg: None
ctx = ScriptRunContext(
"TestSessionID",
fake_enqueue,
"",
SessionState(),
UploadedFileManager(),
)
ctx.on_script_start()
markdown_msg = ForwardMsg()
markdown_msg.delta.new_element.markdown.body = "foo"
msg = ForwardMsg()
msg.page_config_changed.title = "foo"
ctx.enqueue(markdown_msg)
with self.assertRaises(StreamlitAPIException):
ctx.enqueue(msg)
def start(self) -> None:
"""Start a new thread to process the ScriptEventQueue.
This must be called only once.
"""
if self._script_thread is not None:
raise Exception("ScriptRunner was already started")
self._script_thread = threading.Thread(
target=self._process_request_queue,
name="ScriptRunner.scriptThread",
)
script_run_ctx = ScriptRunContext(
session_id=self._session_id,
enqueue=self._enqueue_forward_msg,
query_string=self._client_state.query_string,
session_state=self._session_state,
uploaded_file_mgr=self._uploaded_file_mgr,
)
add_script_run_ctx(self._script_thread, script_run_ctx)
self._script_thread.start()
def test_set_page_config_immutable(self):
"""st.set_page_config must be called at most once"""
fake_enqueue = lambda msg: None
ctx = ScriptRunContext(
"TestSessionID",
fake_enqueue,
"",
SessionState(),
UploadedFileManager(),
)
msg = ForwardMsg()
msg.page_config_changed.title = "foo"
ctx.enqueue(msg)
with self.assertRaises(StreamlitAPIException):
ctx.enqueue(msg)
def test_cached_st_function_warning(self, _, cache_decorator, call_stack):
"""Ensure we properly warn when st.foo functions are called
inside a cached function.
"""
forward_msg_queue = ForwardMsgQueue()
orig_report_ctx = get_script_run_ctx()
add_script_run_ctx(
threading.current_thread(),
ScriptRunContext(
session_id="test session id",
enqueue=forward_msg_queue.enqueue,
query_string="",
session_state=SessionState(),
uploaded_file_mgr=None,
),
)
with patch.object(call_stack,
"_show_cached_st_function_warning") as warning:
st.text("foo")
warning.assert_not_called()
@cache_decorator
def cached_func():
st.text("Inside cached func")
cached_func()
warning.assert_called_once()
warning.reset_mock()
# Make sure everything got reset properly
st.text("foo")
warning.assert_not_called()
# Test warning suppression
@cache_decorator(suppress_st_warning=True)
def suppressed_cached_func():
st.text("No warnings here!")
suppressed_cached_func()
warning.assert_not_called()
# Test nested st.cache functions
@cache_decorator
def outer():
@cache_decorator
def inner():
st.text("Inside nested cached func")
return inner()
outer()
warning.assert_called_once()
warning.reset_mock()
# Test st.cache functions that raise errors
with self.assertRaises(RuntimeError):
@cache_decorator
def cached_raise_error():
st.text("About to throw")
raise RuntimeError("avast!")
cached_raise_error()
warning.assert_called_once()
warning.reset_mock()
# Make sure everything got reset properly
st.text("foo")
warning.assert_not_called()
# Test st.cache functions with widgets
@cache_decorator
def cached_widget():
st.button("Press me!")
cached_widget()
warning.assert_called_once()
warning.reset_mock()
# Make sure everything got reset properly
st.text("foo")
warning.assert_not_called()
add_script_run_ctx(threading.current_thread(), orig_report_ctx)
def test_set_page_config_reset(self):
"""st.set_page_config should be allowed after a rerun"""
fake_enqueue = lambda msg: None
ctx = ScriptRunContext(
"TestSessionID",
fake_enqueue,
"",
SessionState(),
UploadedFileManager(),
)
ctx.on_script_start()
msg = ForwardMsg()
msg.page_config_changed.title = "foo"
ctx.enqueue(msg)
ctx.reset()
try:
ctx.on_script_start()
ctx.enqueue(msg)
except StreamlitAPIException:
self.fail("set_page_config should have succeeded after reset!")
def test_enqueue_new_session_message(self, patched_config):
"""The SCRIPT_STARTED event should enqueue a 'new_session' message."""
def get_option(name):
if name == "server.runOnSave":
# Just to avoid starting the watcher for no reason.
return False
return config.get_option(name)
patched_config.get_option.side_effect = get_option
patched_config.get_options_for_section.side_effect = (
_mock_get_options_for_section()
)
# Create a AppSession with some mocked bits
session = AppSession(
ioloop=self.io_loop,
session_data=SessionData("mock_report.py", ""),
uploaded_file_manager=UploadedFileManager(),
message_enqueued_callback=lambda: None,
local_sources_watcher=MagicMock(),
)
orig_ctx = get_script_run_ctx()
ctx = ScriptRunContext(
session_id="TestSessionID",
enqueue=session._session_data.enqueue,
query_string="",
session_state=MagicMock(),
uploaded_file_mgr=MagicMock(),
)
add_script_run_ctx(ctx=ctx)
# Send a mock SCRIPT_STARTED event.
session._on_scriptrunner_event(
sender=MagicMock(), event=ScriptRunnerEvent.SCRIPT_STARTED
)
sent_messages = session._session_data._browser_queue._queue
self.assertEqual(2, len(sent_messages)) # NewApp and SessionState messages
# Note that we're purposefully not very thoroughly testing new_session
# fields below to avoid getting to the point where we're just
# duplicating code in tests.
new_session_msg = sent_messages[0].new_session
self.assertEqual("mock_scriptrun_id", new_session_msg.script_run_id)
self.assertTrue(new_session_msg.HasField("config"))
self.assertEqual(
config.get_option("server.allowRunOnSave"),
new_session_msg.config.allow_run_on_save,
)
self.assertTrue(new_session_msg.HasField("custom_theme"))
self.assertEqual("black", new_session_msg.custom_theme.text_color)
init_msg = new_session_msg.initialize
self.assertTrue(init_msg.HasField("user_info"))
add_script_run_ctx(ctx=orig_ctx)