Skip to content

Commit

Permalink
[App] Fix environment check with command redirection (#16883)
Browse files Browse the repository at this point in the history
  • Loading branch information
ethanwharris authored and lexierule committed Feb 28, 2023
1 parent d4c5ff3 commit 8e55ff7
Show file tree
Hide file tree
Showing 3 changed files with 74 additions and 8 deletions.
20 changes: 15 additions & 5 deletions src/lightning_app/cli/lightning_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -76,11 +76,21 @@ def main() -> None:
# Check environment and versions if not in the cloud and not testing
is_testing = bool(int(os.getenv("LIGHTING_TESTING", "0")))
if not is_testing and "LIGHTNING_APP_STATE_URL" not in os.environ:
# Enforce running in PATH Python
_check_environment_and_redirect()

# Check for newer versions and upgrade
_check_version_and_upgrade()
try:
# Enforce running in PATH Python
_check_environment_and_redirect()

# Check for newer versions and upgrade
_check_version_and_upgrade()
except SystemExit:
raise
except Exception:
# Note: We intentionally ignore all exceptions here so that we never panic if one of the above calls fails.
# If they fail for some reason users should still be able to continue with their command.
click.echo(
"We encountered an unexpected problem while checking your environment."
"We will still proceed with the command, however, there is a chance that errors may occur."
)

# 1: Handle connection to a Lightning App.
if len(sys.argv) > 1 and sys.argv[1] in ("connect", "disconnect", "logout"):
Expand Down
10 changes: 8 additions & 2 deletions src/lightning_app/utilities/cli_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@
import json
import os
import re
import shutil
import subprocess
import sys
from typing import Dict, Optional
Expand Down Expand Up @@ -323,7 +322,14 @@ def _check_environment_and_redirect():
If not, this utility tries to redirect the ``lightning`` call to the environment executable (prompting the user to
install lightning for them there if needed).
"""
env_executable = os.path.realpath(shutil.which("python"))
process = subprocess.run(
["python", "-c", "import sys; print(sys.executable)"],
capture_output=True,
env=os.environ,
check=True,
)

env_executable = os.path.realpath(process.stdout.decode().strip())
sys_executable = os.path.realpath(sys.executable)

# on windows, the extension might be different, where one uses `.EXE` and the other `.exe`
Expand Down
52 changes: 51 additions & 1 deletion tests/tests_app/utilities/test_cli_helpers.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,17 @@
import os
import sys
from unittest.mock import Mock, patch

import arrow
import pytest

import lightning_app
from lightning_app.utilities.cli_helpers import _arrow_time_callback, _format_input_env_variables, _get_newer_version
from lightning_app.utilities.cli_helpers import (
_arrow_time_callback,
_check_environment_and_redirect,
_format_input_env_variables,
_get_newer_version,
)


def test_format_input_env_variables():
Expand Down Expand Up @@ -111,3 +118,46 @@ def test_get_newer_version(mock_requests, releases, current_version, newer_versi

_get_newer_version.cache_clear()
assert _get_newer_version() == newer_version


@patch("lightning_app.utilities.cli_helpers._redirect_command")
def test_check_environment_and_redirect(mock_redirect_command, tmpdir, monkeypatch):
# Ensure that the test fails if it tries to redirect
mock_redirect_command.side_effect = RuntimeError

# Test normal executable on the path
# Ensure current executable is on the path
monkeypatch.setenv("PATH", f"{os.path.dirname(sys.executable)}")

assert _check_environment_and_redirect() is None

# Test executable on the path with redirect
fake_python_path = os.path.join(tmpdir, "python")

os.symlink(sys.executable, fake_python_path)

monkeypatch.setenv("PATH", f"{tmpdir}")
assert _check_environment_and_redirect() is None

os.remove(fake_python_path)

descriptor = os.open(
fake_python_path,
flags=(
os.O_WRONLY # access mode: write only
| os.O_CREAT # create if not exists
| os.O_TRUNC # truncate the file to zero
),
mode=0o777,
)

with open(descriptor, "w") as f:
f.writelines(
[
"#!/bin/bash\n",
f'{sys.executable} "$@"',
]
)

monkeypatch.setenv("PATH", f"{tmpdir}")
assert _check_environment_and_redirect() is None

0 comments on commit 8e55ff7

Please sign in to comment.