add timeout test

hngynjy · Aug 17, 2021 · f8c7d2b · f8c7d2b
1 parent 9c07e35
commit f8c7d2b
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 36 deletions.
diff --git a/lib/streamlit/server/server.py b/lib/streamlit/server/server.py
@@ -19,7 +19,6 @@
 import sys
 import errno
 import time
-import types
 import traceback
 import click
 from enum import Enum
@@ -34,8 +33,6 @@
 
 from streamlit import config
 from streamlit import file_util
-from streamlit import magic
-from streamlit import source_util
 from streamlit import util
 from streamlit.config_option import ConfigOption
 from streamlit.forward_msg_cache import ForwardMsgCache
@@ -263,7 +260,6 @@ def __init__(
         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]
-        self._script_health_check_session_id = None  # type: Optional[str]
 
     def __repr__(self) -> str:
         return util.repr_(self)
@@ -430,40 +426,35 @@ async def does_script_run_without_error(self) -> Tuple[bool, str]:
 
         Returns
         -------
-        (True, "ok") if the script completes without error, or [False, err_msg]
+        (True, "ok") if the script completes without error, or (False, err_msg)
         if the script raises an exception.
         """
-        if self._script_health_check_session_id is None:
-            session = ReportSession(
-                ioloop=self._ioloop,
-                script_path=self._script_path,
-                command_line=self._command_line,
-                uploaded_file_manager=self._uploaded_file_mgr,
-            )
-            self._script_health_check_session_id = session.id
-            self._session_info_by_id[
-                self._script_health_check_session_id
-            ] = SessionInfo(None, session)
-
-        session = self._session_info_by_id[self._script_health_check_session_id].session
+        session = ReportSession(
+            ioloop=self._ioloop,
+            script_path=self._script_path,
+            command_line=self._command_line,
+            uploaded_file_manager=self._uploaded_file_mgr,
+        )
 
-        session.session_state.clear_state()
-        session.request_rerun(None)
+        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)
+            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"
+            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"
+            ok = session.session_state[SCRIPT_RUN_WITHOUT_ERRORS_KEY]
+            msg = "ok" if ok else "error"
 
-        return ok, msg
+            return ok, msg
+        finally:
+            session.shutdown()
 
     @property
     def browser_is_connected(self):

diff --git a/lib/tests/streamlit/server_test.py b/lib/tests/streamlit/server_test.py
@@ -13,15 +13,12 @@
 # limitations under the License.
 
 """Server.py unit tests"""
-import logging
 import os
 import shutil
-import uuid
 from unittest import mock
-from unittest.mock import MagicMock, patch, mock_open
+from unittest.mock import MagicMock, patch
 import unittest
 import tempfile
-from uuid import UUID
 
 import pytest
 import tornado.testing
@@ -33,7 +30,6 @@
 import streamlit.server.server
 from streamlit import config, RootContainer
 from streamlit.cursor import make_delta_path
-from streamlit.report_session import ReportSession
 from streamlit.uploaded_file_manager import UploadedFileRec
 from streamlit.server.server import MAX_PORT_SEARCH_RETRIES
 from streamlit.forward_msg_cache import ForwardMsgCache
@@ -641,6 +637,16 @@ async def test_valid_script(self):
             "import streamlit as st\n\nst.write('test')", True, "ok"
         )
 
+    @tornado.testing.gen_test(timeout=5)
+    async def test_timeout_script(self):
+        try:
+            streamlit.server.server.SCRIPT_RUN_CHECK_TIMEOUT = 0.1
+            await self._check_script_loading(
+                "import time\n\ntime.sleep(5)", False, "timeout"
+            )
+        finally:
+            streamlit.server.server.SCRIPT_RUN_CHECK_TIMEOUT = 60
+
     async def _check_script_loading(self, script, expected_loads, expected_msg):
         with os.fdopen(self._fd, "w") as tmp:
             tmp.write(script)