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_clear_cache_resets_session_state(self, _1):
rs = AppSession(
None, SessionData("", ""), UploadedFileManager(), None, MagicMock()
)
rs._session_state["foo"] = "bar"
rs.handle_clear_cache_request()
self.assertTrue("foo" not in rs._session_state)
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 __init__(self, ioloop: tornado.ioloop.IOLoop, script_path: str,
command_line: str):
"""Create the server. It won't be started yet."""
if Server._singleton is not None:
raise RuntimeError(
"Server already initialized. Use .get_current() instead")
Server._singleton = self
_set_tornado_log_levels()
self._ioloop = ioloop
self._script_path = script_path
self._command_line = command_line
# Mapping of ReportSession.id -> SessionInfo.
self._session_info_by_id = {} # type: Dict[str, SessionInfo]
self._must_stop = threading.Event()
self._state = None
self._set_state(State.INITIAL)
self._message_cache = ForwardMsgCache()
self._uploaded_file_mgr = UploadedFileManager()
self._uploaded_file_mgr.on_files_updated.connect(self.on_files_updated)
self._report = None # type: Optional[Report]
self._preheated_session_id = None # type: Optional[str]
def __init__(self, script_name):
"""Initializes the ScriptRunner for the given script_name"""
# DeltaGenerator deltas will be enqueued into self.forward_msg_queue.
self.forward_msg_queue = ForwardMsgQueue()
def enqueue_fn(msg):
self.forward_msg_queue.enqueue(msg)
self.maybe_handle_execution_control_request()
self.script_request_queue = ScriptRequestQueue()
script_path = os.path.join(os.path.dirname(__file__), "test_data",
script_name)
super(TestScriptRunner, self).__init__(
session_id="test session id",
session_data=SessionData(script_path, "test command line"),
enqueue_forward_msg=enqueue_fn,
client_state=ClientState(),
session_state=SessionState(),
request_queue=self.script_request_queue,
uploaded_file_mgr=UploadedFileManager(),
)
# Accumulates uncaught exceptions thrown by our run thread.
self.script_thread_exceptions = []
# Accumulates all ScriptRunnerEvents emitted by us.
self.events = []
self.event_data = []
def record_event(event, **kwargs):
self.events.append(event)
self.event_data.append(kwargs)
self.on_event.connect(record_event, weak=False)
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_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 get_app(self):
self.file_mgr = UploadedFileManager()
self._get_session_info = lambda x: True
return tornado.web.Application([
(
UPLOAD_FILE_ROUTE,
UploadFileRequestHandler,
dict(
file_mgr=self.file_mgr,
get_session_info=self._get_session_info,
),
),
])
def test_set_page_config_immutable(self):
"""st.set_page_config must be called at most once"""
fake_enqueue = lambda msg: None
ctx = ReportContext("TestSessionID", fake_enqueue, "", Widgets(),
UploadedFileManager())
msg = ForwardMsg()
msg.page_config_changed.title = "foo"
ctx.enqueue(msg)
with self.assertRaises(StreamlitAPIException):
ctx.enqueue(msg)
def get_app(self):
self.file_mgr = UploadedFileManager()
return tornado.web.Application([
(
"/upload_file/(.*)/(.*)/([0-9]*)?",
UploadFileRequestHandler,
dict(file_mgr=self.file_mgr),
),
(
"/upload_file",
UploadFileRequestHandler,
dict(file_mgr=self.file_mgr),
),
])
def test_set_page_config_first(self):
"""st.set_page_config must be called before other st commands"""
fake_enqueue = lambda msg: None
ctx = ReportContext("TestSessionID", fake_enqueue, "", Widgets(),
UploadedFileManager())
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_reset(self):
"""st.set_page_config should be allowed after a rerun"""
fake_enqueue = lambda msg: None
ctx = ReportContext("TestSessionID", fake_enqueue, "", Widgets(),
UploadedFileManager())
msg = ForwardMsg()
msg.page_config_changed.title = "foo"
ctx.enqueue(msg)
ctx.reset()
try:
ctx.enqueue(msg)
except StreamlitAPIException:
self.fail("set_page_config should have succeeded after reset!")
def setUp(self, override_root=True):
self.report_queue = ReportQueue()
self.override_root = override_root
self.orig_report_ctx = None
if self.override_root:
self.orig_report_ctx = get_report_ctx()
add_report_ctx(
threading.current_thread(),
ReportContext(
session_id="test session id",
enqueue=self.report_queue.enqueue,
query_string="",
widgets=Widgets(),
uploaded_file_mgr=UploadedFileManager(),
),
)
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 __init__(self, script_name: str):
"""Initializes the ScriptRunner for the given script_name"""
# DeltaGenerator deltas will be enqueued into self.forward_msg_queue.
self.forward_msg_queue = ForwardMsgQueue()
main_script_path = os.path.join(os.path.dirname(__file__), "test_data",
script_name)
super().__init__(
session_id="test session id",
session_data=SessionData(main_script_path, "test command line"),
client_state=ClientState(),
session_state=SessionState(),
uploaded_file_mgr=UploadedFileManager(),
initial_rerun_data=RerunData(),
)
# Accumulates uncaught exceptions thrown by our run thread.
self.script_thread_exceptions: List[BaseException] = []
# Accumulates all ScriptRunnerEvents emitted by us.
self.events: List[ScriptRunnerEvent] = []
self.event_data: List[Any] = []
def record_event(sender: Optional[ScriptRunner],
event: ScriptRunnerEvent, **kwargs) -> None:
# Assert that we're not getting unexpected `sender` params
# from ScriptRunner.on_event
assert (sender is None
or sender == self), "Unexpected ScriptRunnerEvent sender!"
self.events.append(event)
self.event_data.append(kwargs)
# 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 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_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 = AppSession(
None, SessionData("", ""), UploadedFileManager(), lambda: None, MagicMock()
)
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 SessionData 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 __init__(self, ioloop: IOLoop, script_path: str,
command_line: Optional[str]):
"""Create the server. It won't be started yet."""
if Server._singleton is not None:
raise RuntimeError(
"Server already initialized. Use .get_current() instead")
Server._singleton = self
_set_tornado_log_levels()
self._ioloop = ioloop
self._script_path = script_path
self._command_line = command_line
# Mapping of ReportSession.id -> SessionInfo.
self._session_info_by_id: Dict[str, SessionInfo] = {}
self._must_stop = tornado.locks.Event()
self._state = State.INITIAL
self._message_cache = ForwardMsgCache()
self._uploaded_file_mgr = UploadedFileManager()
self._uploaded_file_mgr.on_files_updated.connect(self.on_files_updated)
self._report: Optional[Report] = None
self._preheated_session_id: Optional[str] = None
self._has_connection = tornado.locks.Condition()
self._need_send_data = tornado.locks.Event()
# StatsManager
self._stats_mgr = StatsManager()
self._stats_mgr.register_provider(get_memo_stats_provider())
self._stats_mgr.register_provider(get_singleton_stats_provider())
self._stats_mgr.register_provider(_mem_caches)
self._stats_mgr.register_provider(self._message_cache)
self._stats_mgr.register_provider(in_memory_file_manager)
self._stats_mgr.register_provider(self._uploaded_file_mgr)
self._stats_mgr.register_provider(
SessionStateStatProvider(self._session_info_by_id))
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_creates_session_state_on_init(self, _):
rs = AppSession(
None, SessionData("", ""), UploadedFileManager(), None, MagicMock()
)
self.assertTrue(isinstance(rs.session_state, SessionState))
def test_request_rerun_on_secrets_file_change(self, patched_connect):
rs = AppSession(
None, SessionData("", ""), UploadedFileManager(), None, MagicMock()
)
patched_connect.assert_called_once_with(rs._on_secrets_file_changed)
def setUp(self):
self.mgr = UploadedFileManager()
self.filemgr_events = []
self.mgr.on_files_updated.connect(self._on_files_updated)
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)
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))
