def test_enqueue_without_tracer(self, _1, _2, patched_config): """Make sure we try to handle execution control requests.""" def get_option(name): if name == "server.runOnSave": # Just to avoid starting the watcher for no reason. return False if name == "client.displayEnabled": return True if name == "runner.installTracer": return False raise RuntimeError("Unexpected argument to get_option: %s" % name) patched_config.get_option.side_effect = get_option rs = ReportSession(None, "", "", UploadedFileManager()) mock_script_runner = MagicMock() mock_script_runner._install_tracer = ScriptRunner._install_tracer rs._scriptrunner = mock_script_runner mock_msg = MagicMock() rs.enqueue(mock_msg) func = mock_script_runner.maybe_handle_execution_control_request # Expect func to be called only once, inside enqueue(). func.assert_called_once()
def test_clear_cache_all_caches( self, clear_singleton_cache, clear_memo_cache, clear_legacy_cache ): rs = ReportSession(MagicMock(), "", "", UploadedFileManager(), None) rs.handle_clear_cache_request() clear_singleton_cache.assert_called_once() clear_memo_cache.assert_called_once() clear_legacy_cache.assert_called_once()
def test_get_deploy_params_with_no_git(self, _1): """Make sure we try to handle execution control requests.""" import os os.environ["PATH"] = "" rs = ReportSession(None, report_session.__file__, "", UploadedFileManager()) self.assertIsNone(rs.get_deploy_params())
def test_handle_save_request(self, _1): """Test that handle_save_request serializes files correctly.""" # Create a ReportSession with some mocked bits rs = ReportSession( self.io_loop, "mock_report.py", "", UploadedFileManager(), None ) rs._report.report_id = "TestReportID" orig_ctx = get_report_ctx() ctx = ReportContext( "TestSessionID", rs._report.enqueue, "", SessionState(), UploadedFileManager(), ) add_report_ctx(ctx=ctx) rs._scriptrunner = MagicMock() storage = MockStorage() rs._storage = storage # Send two deltas: empty and markdown st.empty() st.markdown("Text!") yield rs.handle_save_request(_create_mock_websocket()) # Check the order of the received files. Manifest should be last. self.assertEqual(3, len(storage.files)) self.assertEqual("reports/TestReportID/0.pb", storage.get_filename(0)) self.assertEqual("reports/TestReportID/1.pb", storage.get_filename(1)) self.assertEqual("reports/TestReportID/manifest.pb", storage.get_filename(2)) # Check the manifest manifest = storage.get_message(2, StaticManifest) self.assertEqual("mock_report", manifest.name) self.assertEqual(2, manifest.num_messages) self.assertEqual(StaticManifest.DONE, manifest.server_status) # Check that the deltas we sent match messages in storage sent_messages = rs._report._master_queue._queue received_messages = [ storage.get_message(0, ForwardMsg), storage.get_message(1, ForwardMsg), ] self.assertEqual(sent_messages, received_messages) add_report_ctx(ctx=orig_ctx)
def test_enqueue_new_report_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 ReportSession with some mocked bits rs = ReportSession(self.io_loop, "mock_report.py", "", UploadedFileManager()) rs._report.report_id = "testing _enqueue_new_report" orig_ctx = get_report_ctx() ctx = ReportContext("TestSessionID", rs._report.enqueue, "", None, None) add_report_ctx(ctx=ctx) rs._on_scriptrunner_event(ScriptRunnerEvent.SCRIPT_STARTED) sent_messages = rs._report._master_queue._queue self.assertEqual(len(sent_messages), 2) # NewReport and SessionState messages # Note that we're purposefully not very thoroughly testing new_report # fields below to avoid getting to the point where we're just # duplicating code in tests. new_report_msg = sent_messages[0].new_report self.assertEqual(new_report_msg.report_id, rs._report.report_id) self.assertEqual(new_report_msg.HasField("config"), True) self.assertEqual( new_report_msg.config.allow_run_on_save, config.get_option("server.allowRunOnSave"), ) self.assertEqual(new_report_msg.HasField("custom_theme"), True) self.assertEqual(new_report_msg.custom_theme.text_color, "black") init_msg = new_report_msg.initialize self.assertEqual(init_msg.HasField("user_info"), True) add_report_ctx(ctx=orig_ctx)
def _create_or_reuse_report_session( self, ws: Optional[WebSocketHandler] ) -> ReportSession: """Register a connected browser with the server. Parameters ---------- ws : _BrowserWebSocketHandler or None The newly-connected websocket handler or None if preheated connection. Returns ------- ReportSession The newly-created ReportSession for this browser connection. """ if self._preheated_session_id is not None: assert len(self._session_info_by_id) == 1 assert ws is not None session_id = self._preheated_session_id self._preheated_session_id = None session_info = self._session_info_by_id[session_id] session_info.ws = ws session = session_info.session LOGGER.debug( "Reused preheated session for ws %s. Session ID: %s", id(ws), session_id ) else: session = ReportSession( ioloop=self._ioloop, script_path=self._script_path, command_line=self._command_line, uploaded_file_manager=self._uploaded_file_mgr, message_enqueued_callback=self._enqueued_some_message, ) LOGGER.debug( "Created new session for ws %s. Session ID: %s", id(ws), session.id ) assert session.id not in self._session_info_by_id, ( "session.id '%s' registered multiple times!" % session.id ) self._session_info_by_id[session.id] = SessionInfo(ws, session) if ws is None: self._preheated_session_id = session.id else: self._set_state(State.ONE_OR_MORE_BROWSERS_CONNECTED) self._has_connection.notify_all() return session
async def does_script_run_without_error(self) -> Tuple[bool, str]: """Load and execute the app's script to verify it runs without an error. Returns ------- (True, "ok") if the script completes without error, or (False, err_msg) if the script raises an exception. """ session = ReportSession( ioloop=self._ioloop, script_path=self._script_path, command_line=self._command_line, uploaded_file_manager=self._uploaded_file_mgr, ) try: session.request_rerun(None) now = time.perf_counter() while ( SCRIPT_RUN_WITHOUT_ERRORS_KEY not in session.session_state and (time.perf_counter() - now) < SCRIPT_RUN_CHECK_TIMEOUT ): await tornado.gen.sleep(0.1) if SCRIPT_RUN_WITHOUT_ERRORS_KEY not in session.session_state: return False, "timeout" ok = session.session_state[SCRIPT_RUN_WITHOUT_ERRORS_KEY] msg = "ok" if ok else "error" return ok, msg finally: session.shutdown()
def test_passes_client_state_on_run_on_save(self, _): rs = ReportSession(None, "", "", UploadedFileManager(), None) rs._run_on_save = True rs.request_rerun = MagicMock() rs._on_source_file_changed() rs.request_rerun.assert_called_once_with(rs._client_state)
def test_enqueue_with_tracer(self, _1, _2, patched_config, _4): """Make sure there is no lock contention when tracer is on. When the tracer is set up, we want maybe_handle_execution_control_request to be executed only once. There was a bug in the past where it was called twice: once from the tracer and once from the enqueue function. This caused a lock contention. """ def get_option(name): if name == "server.runOnSave": # Just to avoid starting the watcher for no reason. return False if name == "client.displayEnabled": return True if name == "runner.installTracer": return True raise RuntimeError("Unexpected argument to get_option: %s" % name) patched_config.get_option.side_effect = get_option rs = ReportSession(None, "", "", UploadedFileManager()) mock_script_runner = MagicMock() rs._scriptrunner = mock_script_runner mock_msg = MagicMock() rs.enqueue(mock_msg) func = mock_script_runner.maybe_handle_execution_control_request # In reality, outside of a testing environment func should be called # once. But in this test we're actually not installing a tracer here, # since Report is mocked. So the correct behavior here is for func to # never be called. If you ever see it being called once here it's # likely because there's a bug in the enqueue function (which should # skip func when installTracer is on). func.assert_not_called()
def test_set_page_config_immutable(self, _1): """st.set_page_config must be called at most once""" file_mgr = MagicMock(spec=UploadedFileManager) rs = ReportSession(None, "", "", file_mgr) msg = ForwardMsg() msg.page_config_changed.title = "foo" rs.enqueue(msg) with self.assertRaises(StreamlitAPIException): rs.enqueue(msg)
def test_shutdown(self, _1): """Test that ReportSession.shutdown behaves sanely.""" file_mgr = MagicMock(spec=UploadedFileManager) rs = ReportSession(None, "", "", file_mgr) rs.shutdown() self.assertEqual(ReportSessionState.SHUTDOWN_REQUESTED, rs._state) file_mgr.remove_session_files.assert_called_once_with(rs.id) # A 2nd shutdown call should have no effect. rs.shutdown() self.assertEqual(ReportSessionState.SHUTDOWN_REQUESTED, rs._state) file_mgr.remove_session_files.assert_called_once_with(rs.id)
def test_set_page_config_first(self, _1): """st.set_page_config must be called before other st commands""" file_mgr = MagicMock(spec=UploadedFileManager) rs = ReportSession(None, "", "", file_mgr) markdown_msg = ForwardMsg() markdown_msg.delta.new_element.markdown.body = "foo" msg = ForwardMsg() msg.page_config_changed.title = "foo" rs.enqueue(markdown_msg) with self.assertRaises(StreamlitAPIException): rs.enqueue(msg)
def test_unique_id(self, _1): """Each ReportSession should have a unique ID""" file_mgr = MagicMock(spec=UploadedFileManager) rs1 = ReportSession(None, "", "", file_mgr) rs2 = ReportSession(None, "", "", file_mgr) self.assertNotEqual(rs1.id, rs2.id)
def test_request_rerun_on_secrets_file_change(self, _, patched_connect): rs = ReportSession(None, "", "", UploadedFileManager(), None) patched_connect.assert_called_once_with(rs._on_secrets_file_changed)
def test_clear_cache_resets_session_state(self, _1): rs = ReportSession(None, "", "", UploadedFileManager(), None) rs._session_state["foo"] = "bar" rs.handle_clear_cache_request() self.assertTrue("foo" not in rs._session_state)
def test_creates_session_state_on_init(self, _): rs = ReportSession(None, "", "", UploadedFileManager(), None) self.assertTrue(isinstance(rs.session_state, SessionState))