def test_enqueue_new_session_message(self):
"""The SCRIPT_STARTED event should enqueue a 'new_session' message."""
session = _create_test_session(self.io_loop)
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)
mock_scriptrunner = MagicMock(spec=ScriptRunner)
session._scriptrunner = mock_scriptrunner
# Send a mock SCRIPT_STARTED event.
session._on_scriptrunner_event(
sender=mock_scriptrunner,
event=ScriptRunnerEvent.SCRIPT_STARTED,
)
# Yield to let the AppSession's callbacks run.
yield
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)
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 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_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_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)