From 8e55ff753dc7cdff8469f29472e24ea82fac41e6 Mon Sep 17 00:00:00 2001
From: Ethan Harris <ethanwharris@gmail.com>
Date: Tue, 28 Feb 2023 12:34:43 +0000
Subject: [PATCH] [App] Fix environment check with command redirection (#16883)

---
 src/lightning_app/cli/lightning_cli.py        | 20 +++++--
 src/lightning_app/utilities/cli_helpers.py    | 10 +++-
 tests/tests_app/utilities/test_cli_helpers.py | 52 ++++++++++++++++++-
 3 files changed, 74 insertions(+), 8 deletions(-)

diff --git a/src/lightning_app/cli/lightning_cli.py b/src/lightning_app/cli/lightning_cli.py
index 4ff505c42ae12..bfb0fff543f21 100644
--- a/src/lightning_app/cli/lightning_cli.py
+++ b/src/lightning_app/cli/lightning_cli.py
@@ -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"):
diff --git a/src/lightning_app/utilities/cli_helpers.py b/src/lightning_app/utilities/cli_helpers.py
index 701da11bad533..56b76e076832b 100644
--- a/src/lightning_app/utilities/cli_helpers.py
+++ b/src/lightning_app/utilities/cli_helpers.py
@@ -16,7 +16,6 @@
 import json
 import os
 import re
-import shutil
 import subprocess
 import sys
 from typing import Dict, Optional
@@ -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`
diff --git a/tests/tests_app/utilities/test_cli_helpers.py b/tests/tests_app/utilities/test_cli_helpers.py
index ecdd1705c2130..a6a80b4ad3192 100644
--- a/tests/tests_app/utilities/test_cli_helpers.py
+++ b/tests/tests_app/utilities/test_cli_helpers.py
@@ -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():
@@ -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