streamlit

1
#!/usr/bin/env python
2

3
# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2024)
4
#
5
# Licensed under the Apache License, Version 2.0 (the "License");
6
# you may not use this file except in compliance with the License.
7
# You may obtain a copy of the License at
8
#
9
#     http://www.apache.org/licenses/LICENSE-2.0
10
#
11
# Unless required by applicable law or agreed to in writing, software
12
# distributed under the License is distributed on an "AS IS" BASIS,
13
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
# See the License for the specific language governing permissions and
15
# limitations under the License.
16

17
import os
18
import shutil
19
import signal
20
import subprocess
21
import sys
22
import time
23
from contextlib import contextmanager
24
from os.path import abspath, basename, dirname, join, splitext
25
from pathlib import Path
26
from tempfile import TemporaryFile
27
from typing import List
28

29
import click
30
import requests
31

32
ROOT_DIR = dirname(dirname(abspath(__file__)))  # streamlit root directory
33
FRONTEND_DIR = join(ROOT_DIR, "frontend")
34

35
CREDENTIALS_FILE = os.path.expanduser("~/.streamlit/credentials.toml")
36

37

38
class QuitException(BaseException):
39
    pass
40

41

42
class AsyncSubprocess:
43
    """A context manager. Wraps subprocess.Popen to capture output safely."""
44

45
    def __init__(self, args, cwd=None, env=None):
46
        self.args = args
47
        self.cwd = cwd
48
        self.env = env
49
        self._proc = None
50
        self._stdout_file = None
51

52
    def terminate(self):
53
        """Terminate the process and return its stdout/stderr in a string."""
54
        # Terminate the process
55
        if self._proc is not None:
56
            self._proc.terminate()
57
            self._proc.wait()
58
            self._proc = None
59

60
        # Read the stdout file and close it
61
        stdout = None
62
        if self._stdout_file is not None:
63
            self._stdout_file.seek(0)
64
            stdout = self._stdout_file.read()
65
            self._stdout_file.close()
66
            self._stdout_file = None
67

68
        return stdout
69

70
    def __enter__(self):
71
        self.start()
72
        return self
73

74
    def start(self):
75
        # Start the process and capture its stdout/stderr output to a temp
76
        # file. We do this instead of using subprocess.PIPE (which causes the
77
        # Popen object to capture the output to its own internal buffer),
78
        # because large amounts of output can cause it to deadlock.
79
        self._stdout_file = TemporaryFile("w+")
80
        self._proc = subprocess.Popen(
81
            self.args,
82
            cwd=self.cwd,
83
            stdout=self._stdout_file,
84
            stderr=subprocess.STDOUT,
85
            text=True,
86
            env={**os.environ.copy(), **self.env} if self.env else None,
87
        )
88

89
    def __exit__(self, exc_type, exc_val, exc_tb):
90
        if self._proc is not None:
91
            self._proc.terminate()
92
            self._proc = None
93
        if self._stdout_file is not None:
94
            self._stdout_file.close()
95
            self._stdout_file = None
96

97

98
class Context:
99
    def __init__(self):
100
        # Whether to prompt to continue on failure or run all
101
        self.always_continue = False
102
        # True if Cypress will record videos of our results.
103
        self.record_results = False
104
        # True if we're automatically updating snapshots.
105
        self.update_snapshots = False
106
        # Parent folder of the specs and scripts.
107
        # 'e2e' for tests we expect to pass or 'e2e_flaky' for tests with
108
        # known issues.
109
        self.tests_dir_name = "e2e"
110
        # Set to True if any test fails.
111
        self.any_failed = False
112
        # Environment variables to pass to Cypress
113
        self.cypress_env_vars = {}
114

115
    @property
116
    def tests_dir(self) -> str:
117
        return join(ROOT_DIR, self.tests_dir_name)
118

119
    @property
120
    def cypress_flags(self) -> List[str]:
121
        """Flags to pass to Cypress"""
122
        flags = ["--config", f"integrationFolder={self.tests_dir}/specs"]
123
        if self.record_results:
124
            flags.append("--record")
125
        if self.update_snapshots:
126
            flags.extend(["--env", "updateSnapshots=true"])
127
        if self.cypress_env_vars:
128
            vars_str = ",".join(f"{k}={v}" for k, v in self.cypress_env_vars.items())
129
            flags.extend(["--env", vars_str])
130
        return flags
131

132

133
def remove_if_exists(path):
134
    """Remove the given folder or file if it exists"""
135
    if os.path.isfile(path):
136
        os.remove(path)
137
    elif os.path.isdir(path):
138
        shutil.rmtree(path)
139

140

141
@contextmanager
142
def move_aside_file(path):
143
    """Move a file aside if it exists; restore it on completion"""
144
    moved = False
145
    if os.path.exists(path):
146
        os.rename(path, f"{path}.bak")
147
        moved = True
148

149
    try:
150
        yield None
151
    finally:
152
        if moved:
153
            os.rename(f"{path}.bak", path)
154

155

156
def create_credentials_toml(contents):
157
    """Writes ~/.streamlit/credentials.toml"""
158
    os.makedirs(dirname(CREDENTIALS_FILE), exist_ok=True)
159
    with open(CREDENTIALS_FILE, "w") as f:
160
        f.write(contents)
161

162

163
def kill_with_pgrep(search_string):
164
    result = subprocess.run(
165
        f"pgrep -f '{search_string}'",
166
        shell=True,
167
        universal_newlines=True,
168
        capture_output=True,
169
    )
170

171
    if result.returncode == 0:
172
        for pid in result.stdout.split():
173
            try:
174
                os.kill(int(pid), signal.SIGTERM)
175
            except Exception as e:
176
                print("Failed to kill process", e)
177

178

179
def kill_streamlits():
180
    """Kill any active `streamlit run` processes"""
181
    kill_with_pgrep("streamlit run")
182

183

184
def kill_app_servers():
185
    """Kill any active app servers spawned by this script."""
186
    kill_with_pgrep("running-streamlit-e2e-test")
187

188

189
def run_test(
190
    ctx: Context,
191
    specpath: str,
192
    streamlit_command: List[str],
193
    show_output: bool = False,
194
) -> bool:
195
    """Run a single e2e test.
196

197
     An e2e test consists of a Streamlit script that produces a result, and
198
     a Cypress test file that asserts that result is as expected.
199

200
    Parameters
201
    ----------
202
    ctx : Context
203
        The Context object that contains our global testing parameters.
204
    specpath : str
205
        The path of the Cypress spec file to run.
206
    streamlit_command : list of str
207
        The Streamlit command to run (passed directly to subprocess.Popen()).
208

209
    Returns
210
    -------
211
    bool
212
        True if the test succeeded.
213

214
    """
215
    SUCCESS = "SUCCESS"
216
    RETRY = "RETRY"
217
    SKIP = "SKIP"
218
    QUIT = "QUIT"
219

220
    result = None
221

222
    # Move existing credentials file aside, and create a new one if the
223
    # tests call for it.
224
    with move_aside_file(CREDENTIALS_FILE):
225
        create_credentials_toml('[general]\nemail="test@streamlit.io"')
226

227
        # Loop until the test succeeds or is skipped.
228
        while result not in (SUCCESS, SKIP, QUIT):
229
            cypress_command = ["yarn", "cy:run", "--spec", specpath]
230
            cypress_command.extend(ctx.cypress_flags)
231

232
            click.echo(
233
                f"{click.style('Running test:', fg='yellow', bold=True)}"
234
                f"\n{click.style(' '.join(streamlit_command), fg='yellow')}"
235
                f"\n{click.style(' '.join(cypress_command), fg='yellow')}"
236
            )
237

238
            # Start the streamlit command
239
            with AsyncSubprocess(streamlit_command, cwd=FRONTEND_DIR) as streamlit_proc:
240
                # Run the Cypress spec to completion.
241
                cypress_result = subprocess.run(
242
                    cypress_command,
243
                    cwd=FRONTEND_DIR,
244
                    capture_output=True,
245
                    text=True,
246
                )
247

248
                # Terminate the streamlit command and get its output
249
                streamlit_stdout = streamlit_proc.terminate()
250

251
            def print_output():
252
                click.echo(
253
                    f"\n\n{click.style('Streamlit output:', fg='yellow', bold=True)}"
254
                    f"\n{streamlit_stdout}"
255
                    f"\n\n{click.style('Cypress output:', fg='yellow', bold=True)}"
256
                    f"\n{cypress_result.stdout}"
257
                    f"\n"
258
                )
259

260
            if cypress_result.returncode == 0:
261
                result = SUCCESS
262
                click.echo(click.style("Success!\n", fg="green", bold=True))
263
                if show_output:
264
                    print_output()
265
            else:
266
                # The test failed. Print the output of the Streamlit command
267
                # and the Cypress command.
268
                click.echo(click.style("Failure!", fg="red", bold=True))
269
                print_output()
270

271
                if ctx.always_continue:
272
                    result = SKIP
273
                else:
274
                    # Prompt the user for what to do next.
275
                    user_input = click.prompt(
276
                        "[R]etry, [U]pdate snapshots, [S]kip, or [Q]uit?",
277
                        default="r",
278
                    )
279
                    key = user_input[0].lower()
280
                    if key == "s":
281
                        result = SKIP
282
                    elif key == "q":
283
                        result = QUIT
284
                    elif key == "r":
285
                        result = RETRY
286
                    elif key == "u":
287
                        ctx.update_snapshots = True
288
                        result = RETRY
289
                    else:
290
                        # Retry if key not recognized
291
                        result = RETRY
292

293
    if result != SUCCESS:
294
        ctx.any_failed = True
295

296
    if result == QUIT:
297
        raise QuitException()
298

299
    return result == SUCCESS
300

301

302
def is_app_server_alive():
303
    try:
304
        r = requests.get("http://localhost:3000/", timeout=3)
305
        return r.status_code == requests.codes.ok
306
    except:
307
        return False
308

309

310
def run_app_server():
311
    if is_app_server_alive():
312
        print("Detected React app server already running, won't spawn a new one.")
313
        return
314

315
    env = {
316
        "BROWSER": "none",  # don't open up chrome, streamlit does this for us
317
        "BUILD_AS_FAST_AS_POSSIBLE": "true",
318
        "GENERATE_SOURCEMAP": "false",
319
        "INLINE_RUNTIME_CHUNK": "false",
320
    }
321
    command = ["yarn", "start", "--running-streamlit-e2e-test"]
322
    proc = AsyncSubprocess(command, cwd=FRONTEND_DIR, env=env)
323

324
    print("Starting React app server...")
325
    proc.start()
326

327
    print("Waiting for React app server to come online...")
328
    start_time = time.time()
329
    while not is_app_server_alive():
330
        time.sleep(3)
331

332
        # after 10 minutes, we have a problem, just exit
333
        if time.time() - start_time > 60 * 10:
334
            print(
335
                "React app server seems to have had difficulty starting, exiting. Output:"
336
            )
337
            print(proc.terminate())
338
            sys.exit(1)
339

340
    print("React app server is alive!")
341
    return proc
342

343

344
@click.command(
345
    help=(
346
        "Run Streamlit e2e tests. If specific tests are specified, only those "
347
        "tests will be run. If you don't specify specific tests, all tests "
348
        "will be run."
349
    )
350
)
351
@click.option(
352
    "-a", "--always-continue", is_flag=True, help="Continue running on test failure."
353
)
354
@click.option(
355
    "-r",
356
    "--record-results",
357
    is_flag=True,
358
    help="Upload video results to the Cypress dashboard. "
359
    "See https://docs.cypress.io/guides/dashboard/introduction.html for more details.",
360
)
361
@click.option(
362
    "-u",
363
    "--update-snapshots",
364
    is_flag=True,
365
    help="Automatically update snapshots for failing tests.",
366
)
367
@click.option(
368
    "-f",
369
    "--flaky-tests",
370
    is_flag=True,
371
    help="Run tests in 'e2e_flaky' instead of 'e2e'.",
372
)
373
@click.option(
374
    "-v",
375
    "--verbose",
376
    is_flag=True,
377
    help="Show Streamlit and Cypress output.",
378
)
379
@click.argument("tests", nargs=-1)
380
def run_e2e_tests(
381
    always_continue: bool,
382
    record_results: bool,
383
    update_snapshots: bool,
384
    flaky_tests: bool,
385
    tests: List[str],
386
    verbose: bool,
387
):
388
    """Run e2e tests. If any fail, exit with non-zero status."""
389
    kill_streamlits()
390
    kill_app_servers()
391
    app_server = run_app_server()
392

393
    # Clear reports from previous runs
394
    remove_if_exists("frontend/test_results/cypress")
395

396
    ctx = Context()
397
    ctx.always_continue = always_continue
398
    ctx.record_results = record_results
399
    ctx.update_snapshots = update_snapshots
400
    ctx.tests_dir_name = "e2e_flaky" if flaky_tests else "e2e"
401

402
    try:
403
        p = Path(join(ROOT_DIR, ctx.tests_dir_name, "specs")).resolve()
404
        if tests:
405
            paths = [Path(t).resolve() for t in tests]
406
        else:
407
            paths = sorted(p.glob("*.spec.js"))
408
        for spec_path in paths:
409
            if basename(spec_path) == "st_hello.spec.js":
410
                if flaky_tests:
411
                    continue
412

413
                # Test "streamlit hello" in both headless and non-headless mode.
414
                run_test(
415
                    ctx,
416
                    str(spec_path),
417
                    ["streamlit", "hello", "--server.headless=false"],
418
                    show_output=verbose,
419
                )
420
                run_test(
421
                    ctx,
422
                    str(spec_path),
423
                    ["streamlit", "hello", "--server.headless=true"],
424
                    show_output=verbose,
425
                )
426

427
            elif basename(spec_path) == "staticfiles_app.spec.js":
428
                test_name, _ = splitext(basename(spec_path))
429
                test_name, _ = splitext(test_name)
430
                test_path = join(
431
                    ctx.tests_dir,
432
                    "scripts",
433
                    "staticfiles_apps",
434
                    "streamlit_static_app.py",
435
                )
436
                if os.path.exists(test_path):
437
                    run_test(
438
                        ctx,
439
                        str(spec_path),
440
                        [
441
                            "streamlit",
442
                            "run",
443
                            "--server.enableStaticServing=true",
444
                            test_path,
445
                        ],
446
                        show_output=verbose,
447
                    )
448
            elif basename(spec_path) == "hostframe.spec.js":
449
                test_name, _ = splitext(basename(spec_path))
450
                test_name, _ = splitext(test_name)
451
                test_path = join(
452
                    ctx.tests_dir, "scripts", "hostframe", "hostframe_app.py"
453
                )
454
                if os.path.exists(test_path):
455
                    run_test(
456
                        ctx,
457
                        str(spec_path),
458
                        [
459
                            "streamlit",
460
                            "run",
461
                            test_path,
462
                        ],
463
                        show_output=verbose,
464
                    )
465

466
            else:
467
                test_name, _ = splitext(basename(spec_path))
468
                test_name, _ = splitext(test_name)
469
                test_path = join(ctx.tests_dir, "scripts", f"{test_name}.py")
470
                if os.path.exists(test_path):
471
                    run_test(
472
                        ctx,
473
                        str(spec_path),
474
                        ["streamlit", "run", test_path],
475
                        show_output=verbose,
476
                    )
477
    except QuitException:
478
        # Swallow the exception we raise if the user chooses to exit early.
479
        pass
480
    finally:
481
        if app_server:
482
            app_server.terminate()
483

484
    if ctx.any_failed:
485
        sys.exit(1)
486

487

488
if __name__ == "__main__":
489
    run_e2e_tests()
490