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)