Merge pull request #71 from streamlit/STREAMLIT-4007-script-to-check-if-e2e-setup-is-identical

Author: sfc-gh-pbelczyk

.github/workflows/ci.yaml

@@ -51,6 +51,9 @@ jobs:
     - name: Check dependencies for examples
       run: ./dev.py examples-check-deps
 
+    - name: Check e2e utils files
+      run: ./dev.py e2e-utils-check
+
     - name: Checkout streamlit/streamlit
       if: matrix.component_lib_version == 'develop'
       uses: actions/checkout@v3

cookiecutter/{{ cookiecutter.package_name }}/e2e/e2e_utils.py

@@ -17,6 +17,7 @@
 
 
 def _find_free_port():
+    """Find and return a free port on the local machine."""
     with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
         s.bind(("", 0))  # 0 means that the OS chooses a random port
         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
@@ -26,14 +27,22 @@ def _find_free_port():
 class AsyncSubprocess:
     """A context manager. Wraps subprocess. Popen to capture output safely."""
 
-    def __init__(self, args, cwd=None, env=None):
+    def __init__(self, args: typing.List[str], cwd: typing.Optional[str] = None,
+                 env: typing.Optional[typing.Dict[str, str]] = None):
+        """Initialize an AsyncSubprocess instance.
+
+        Args:
+            args (List[str]): List of command-line arguments.
+            cwd (str, optional): Current working directory. Defaults to None.
+            env (dict, optional): Environment variables. Defaults to None.
+        """
         self.args = args
         self.cwd = cwd
         self.env = env
         self._proc = None
         self._stdout_file = None
 
-    def terminate(self):
+    def terminate(self) -> typing.Optional[str]:
         """Terminate the process and return its stdout/stderr in a string."""
         if self._proc is not None:
             self._proc.terminate()
@@ -50,11 +59,13 @@ def terminate(self):
 
         return stdout
 
-    def __enter__(self):
+    def __enter__(self) -> "AsyncSubprocess":
+        """Start the subprocess when entering the context."""
         self.start()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        """Stop the subprocess and close resources when exiting the context."""
         self.stop()
 
     def start(self):
@@ -74,6 +85,7 @@ def start(self):
         )
 
     def stop(self):
+        """Terminate the subprocess and close resources."""
         if self._proc is not None:
             self._proc.terminate()
             self._proc = None
@@ -83,21 +95,32 @@ def stop(self):
 
 
 class StreamlitRunner:
+    """A context manager for running Streamlit scripts."""
+
     def __init__(
             self, script_path: os.PathLike, server_port: typing.Optional[int] = None
     ):
+        """Initialize a StreamlitRunner instance.
+
+        Args:
+            script_path (os.PathLike): Path to the Streamlit script to run.
+            server_port (int, optional): Port for the Streamlit server. Defaults to None.
+        """
         self._process = None
         self.server_port = server_port
         self.script_path = script_path
 
-    def __enter__(self):
+    def __enter__(self) -> "StreamlitRunner":
+        """Start the Streamlit server when entering the context."""
         self.start()
         return self
 
     def __exit__(self, type, value, traceback):
+        """Stop the Streamlit server and close resources when exiting the context."""
         self.stop()
 
     def start(self):
+        """Start the Streamlit server using the specified script and options."""
         self.server_port = self.server_port or _find_free_port()
         self._process = AsyncSubprocess(
             [
@@ -118,12 +141,20 @@ def start(self):
             raise RuntimeError("Application failed to start")
 
     def stop(self):
+        """Stop the Streamlit server and close resources."""
         self._process.stop()
 
-    def is_server_running(self, timeout: int = 30):
+    def is_server_running(self, timeout: int = 30) -> bool:
+        """Check if the Streamlit server is running.
+
+        Args:
+            timeout (int, optional): Maximum time to wait for the server to start. Defaults to 30.
+
+        Returns:
+            bool: True if the server is running, False otherwise.
+        """
         with requests.Session() as http_session:
             start_time = time.time()
-            print("Start loop: ", start_time)
             while True:
                 with contextlib.suppress(requests.RequestException):
                     response = http_session.get(self.server_url + "/_stcore/health")
@@ -134,7 +165,8 @@ def is_server_running(self, timeout: int = 30):
                     return False
 
     @property
-    def server_url(self):
+    def server_url(self) -> str:
+        """Get the URL of the Streamlit server."""
         if not self.server_port:
             raise RuntimeError("Unknown server port")
         return f"http://localhost:{self.server_port}"

dev.py

@@ -6,6 +6,7 @@
 To list the available commands, run ./dev.py --help.
 """
 import argparse
+import glob
 import tempfile
 import typing
 import shlex
@@ -169,6 +170,26 @@ def cmd_example_check_deps(args):
     sys.exit(exit_code)
 
 
+def cmd_check_test_utils(args):
+    """Check that e2e utils files are identical"""
+    file_list = glob.glob('**/e2e_utils.py', recursive=True)
+    if file_list:
+        reference_file = file_list[0]
+    else:
+        print("Cannot find e2e_utils.py files")
+        sys.exit(1)
+
+    for file_path in file_list:
+        run_verbose([
+            "git",
+            "--no-pager",
+            "diff",
+            "--no-index",
+            str(reference_file),
+            str(file_path),
+        ])
+
+
 class CookiecutterVariant(typing.NamedTuple):
     replay_file: Path
     repo_directory: Path
@@ -270,6 +291,7 @@ def cmd_update_templates(args):
     "examples-check-deps": cmd_example_check_deps,
     "templates-check-not-modified": cmd_check_templates_using_cookiecutter,
     "templates-update": cmd_update_templates,
+    "e2e-utils-check": cmd_check_test_utils,
     "e2e-build-images": cmd_e2e_build_images,
     "e2e-run-tests": cmd_e2e_run,
     "docker-images-cleanup": cmd_docker_images_cleanup

examples/CustomDataframe/e2e/e2e_utils.py

@@ -17,6 +17,7 @@
 
 
 def _find_free_port():
+    """Find and return a free port on the local machine."""
     with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
         s.bind(("", 0))  # 0 means that the OS chooses a random port
         s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
@@ -26,14 +27,22 @@ def _find_free_port():
 class AsyncSubprocess:
     """A context manager. Wraps subprocess. Popen to capture output safely."""
 
-    def __init__(self, args, cwd=None, env=None):
+    def __init__(self, args: typing.List[str], cwd: typing.Optional[str] = None,
+                 env: typing.Optional[typing.Dict[str, str]] = None):
+        """Initialize an AsyncSubprocess instance.
+
+        Args:
+            args (List[str]): List of command-line arguments.
+            cwd (str, optional): Current working directory. Defaults to None.
+            env (dict, optional): Environment variables. Defaults to None.
+        """
         self.args = args
         self.cwd = cwd
         self.env = env
         self._proc = None
         self._stdout_file = None
 
-    def terminate(self):
+    def terminate(self) -> typing.Optional[str]:
         """Terminate the process and return its stdout/stderr in a string."""
         if self._proc is not None:
             self._proc.terminate()
@@ -50,11 +59,13 @@ def terminate(self):
 
         return stdout
 
-    def __enter__(self):
+    def __enter__(self) -> "AsyncSubprocess":
+        """Start the subprocess when entering the context."""
         self.start()
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
+        """Stop the subprocess and close resources when exiting the context."""
         self.stop()
 
     def start(self):
@@ -74,6 +85,7 @@ def start(self):
         )
 
     def stop(self):
+        """Terminate the subprocess and close resources."""
         if self._proc is not None:
             self._proc.terminate()
             self._proc = None
@@ -83,21 +95,32 @@ def stop(self):
 
 
 class StreamlitRunner:
+    """A context manager for running Streamlit scripts."""
+
     def __init__(
             self, script_path: os.PathLike, server_port: typing.Optional[int] = None
     ):
+        """Initialize a StreamlitRunner instance.
+
+        Args:
+            script_path (os.PathLike): Path to the Streamlit script to run.
+            server_port (int, optional): Port for the Streamlit server. Defaults to None.
+        """
         self._process = None
         self.server_port = server_port
         self.script_path = script_path
 
-    def __enter__(self):
+    def __enter__(self) -> "StreamlitRunner":
+        """Start the Streamlit server when entering the context."""
         self.start()
         return self
 
     def __exit__(self, type, value, traceback):
+        """Stop the Streamlit server and close resources when exiting the context."""
         self.stop()
 
     def start(self):
+        """Start the Streamlit server using the specified script and options."""
         self.server_port = self.server_port or _find_free_port()
         self._process = AsyncSubprocess(
             [
@@ -118,12 +141,20 @@ def start(self):
             raise RuntimeError("Application failed to start")
 
     def stop(self):
+        """Stop the Streamlit server and close resources."""
         self._process.stop()
 
-    def is_server_running(self, timeout: int = 30):
+    def is_server_running(self, timeout: int = 30) -> bool:
+        """Check if the Streamlit server is running.
+
+        Args:
+            timeout (int, optional): Maximum time to wait for the server to start. Defaults to 30.
+
+        Returns:
+            bool: True if the server is running, False otherwise.
+        """
         with requests.Session() as http_session:
             start_time = time.time()
-            print("Start loop: ", start_time)
             while True:
                 with contextlib.suppress(requests.RequestException):
                     response = http_session.get(self.server_url + "/_stcore/health")
@@ -134,7 +165,8 @@ def is_server_running(self, timeout: int = 30):
                     return False
 
     @property
-    def server_url(self):
+    def server_url(self) -> str:
+        """Get the URL of the Streamlit server."""
         if not self.server_port:
             raise RuntimeError("Unknown server port")
         return f"http://localhost:{self.server_port}"