Win: fix running tests in a virtual environment (#2216)

On windows, starting with python 3.7,
virtual environments use a venvlauncher startup process
This does not play well when counting spawned processes or
when relying on the pid of the spawned process to do some checks
e.g. connection check per pid

This commit detects this situation and uses the base python
executable to spawn processes when required.

Signed-off-by: mayeut <mayeut@users.noreply.github.com>

4 files changed, 30 insertions, 18 deletions

diff --git a/psutil/tests/__init__.py b/psutil/tests/__init__.py
index 61a06e21..e3766eea 100644
--- a/psutil/tests/__init__.py
+++ b/psutil/tests/__init__.py
@@ -41,7 +41,6 @@ from socket import SOCK_STREAM
 
 import psutil
 from psutil import AIX
-from psutil import FREEBSD
 from psutil import LINUX
 from psutil import MACOS
 from psutil import POSIX
@@ -80,8 +79,8 @@ if POSIX:
 __all__ = [
     # constants
     'APPVEYOR', 'DEVNULL', 'GLOBAL_TIMEOUT', 'TOLERANCE_SYS_MEM', 'NO_RETRIES',
-    'PYPY', 'PYTHON_EXE', 'ROOT_DIR', 'SCRIPTS_DIR', 'TESTFN_PREFIX',
-    'UNICODE_SUFFIX', 'INVALID_UNICODE_SUFFIX',
+    'PYPY', 'PYTHON_EXE', 'PYTHON_EXE_ENV', 'ROOT_DIR', 'SCRIPTS_DIR',
+    'TESTFN_PREFIX', 'UNICODE_SUFFIX', 'INVALID_UNICODE_SUFFIX',
     'CI_TESTING', 'VALID_PROC_STATUSES', 'TOLERANCE_DISK_USAGE', 'IS_64BIT',
     "HAS_CPU_AFFINITY", "HAS_CPU_FREQ", "HAS_ENVIRON", "HAS_PROC_IO_COUNTERS",
     "HAS_IONICE", "HAS_MEMORY_MAPS", "HAS_PROC_CPU_NUM", "HAS_RLIMIT",
@@ -240,13 +239,21 @@ def _get_py_exe():
         else:
             return exe
 
-    if GITHUB_ACTIONS:
-        if PYPY:
-            return which("pypy3") if PY3 else which("pypy")
-        elif FREEBSD:
-            return os.path.realpath(sys.executable)
-        else:
-            return which('python')
+    env = os.environ.copy()
+
+    # On Windows, starting with python 3.7, virtual environments use a
+    # venv launcher startup process. This does not play well when
+    # counting spawned processes, or when relying on the PID of the
+    # spawned process to do some checks, e.g. connections check per PID.
+    # Let's use the base python in this case.
+    base = getattr(sys, "_base_executable", None)
+    if WINDOWS and sys.version_info >= (3, 7) and base is not None:
+        # We need to set __PYVENV_LAUNCHER__ to sys.executable for the 
+        # base python executable to know about the environment.
+        env["__PYVENV_LAUNCHER__"] = sys.executable
+        return base, env
+    elif GITHUB_ACTIONS:
+        return sys.executable, env
     elif MACOS:
         exe = \
             attempt(sys.executable) or \
@@ -255,14 +262,14 @@ def _get_py_exe():
             attempt(psutil.Process().exe())
         if not exe:
             raise ValueError("can't find python exe real abspath")
-        return exe
+        return exe, env
     else:
         exe = os.path.realpath(sys.executable)
         assert os.path.exists(exe), exe
-        return exe
+        return exe, env
 
 
-PYTHON_EXE = _get_py_exe()
+PYTHON_EXE, PYTHON_EXE_ENV = _get_py_exe()
 DEVNULL = open(os.devnull, 'r+')
 atexit.register(DEVNULL.close)
 
@@ -351,7 +358,7 @@ def spawn_testproc(cmd=None, **kwds):
     kwds.setdefault("stdin", DEVNULL)
     kwds.setdefault("stdout", DEVNULL)
     kwds.setdefault("cwd", os.getcwd())
-    kwds.setdefault("env", os.environ)
+    kwds.setdefault("env", PYTHON_EXE_ENV)
     if WINDOWS:
         # Prevents the subprocess to open error dialogs. This will also
         # cause stderr to be suppressed, which is suboptimal in order
diff --git a/psutil/tests/test_misc.py b/psutil/tests/test_misc.py
index 44922216..afa60b1c 100755
--- a/psutil/tests/test_misc.py
+++ b/psutil/tests/test_misc.py
@@ -45,6 +45,7 @@ from psutil.tests import HAS_SENSORS_BATTERY
 from psutil.tests import HAS_SENSORS_FANS
 from psutil.tests import HAS_SENSORS_TEMPERATURES
 from psutil.tests import PYTHON_EXE
+from psutil.tests import PYTHON_EXE_ENV
 from psutil.tests import SCRIPTS_DIR
 from psutil.tests import PsutilTestCase
 from psutil.tests import mock
@@ -820,6 +821,7 @@ class TestScripts(PsutilTestCase):
 
     @staticmethod
     def assert_stdout(exe, *args, **kwargs):
+        kwargs.setdefault("env", PYTHON_EXE_ENV)
         exe = '%s' % os.path.join(SCRIPTS_DIR, exe)
         cmd = [PYTHON_EXE, exe]
         for arg in args:
diff --git a/psutil/tests/test_process.py b/psutil/tests/test_process.py
index 26869e98..df6a84c3 100755
--- a/psutil/tests/test_process.py
+++ b/psutil/tests/test_process.py
@@ -52,6 +52,7 @@ from psutil.tests import HAS_THREADS
 from psutil.tests import MACOS_11PLUS
 from psutil.tests import PYPY
 from psutil.tests import PYTHON_EXE
+from psutil.tests import PYTHON_EXE_ENV
 from psutil.tests import PsutilTestCase
 from psutil.tests import ThreadTask
 from psutil.tests import call_until
@@ -1543,7 +1544,7 @@ class TestPopen(PsutilTestCase):
         # Not sure what to do though.
         cmd = [PYTHON_EXE, "-c", "import time; time.sleep(60);"]
         with psutil.Popen(cmd, stdout=subprocess.PIPE,
-                          stderr=subprocess.PIPE) as proc:
+                          stderr=subprocess.PIPE, env=PYTHON_EXE_ENV) as proc:
             proc.name()
             proc.cpu_times()
             proc.stdin
@@ -1559,7 +1560,7 @@ class TestPopen(PsutilTestCase):
         with psutil.Popen([PYTHON_EXE, "-V"],
                           stdout=subprocess.PIPE,
                           stderr=subprocess.PIPE,
-                          stdin=subprocess.PIPE) as proc:
+                          stdin=subprocess.PIPE, env=PYTHON_EXE_ENV) as proc:
             proc.communicate()
         assert proc.stdout.closed
         assert proc.stderr.closed
@@ -1572,7 +1573,7 @@ class TestPopen(PsutilTestCase):
         # diverges from that.
         cmd = [PYTHON_EXE, "-c", "import time; time.sleep(60);"]
         with psutil.Popen(cmd, stdout=subprocess.PIPE,
-                          stderr=subprocess.PIPE) as proc:
+                          stderr=subprocess.PIPE, env=PYTHON_EXE_ENV) as proc:
             proc.terminate()
             proc.wait()
             self.assertRaises(psutil.NoSuchProcess, proc.terminate)
diff --git a/psutil/tests/test_testutils.py b/psutil/tests/test_testutils.py
index 0298ea4e..e757e017 100755
--- a/psutil/tests/test_testutils.py
+++ b/psutil/tests/test_testutils.py
@@ -30,6 +30,7 @@ from psutil.tests import CI_TESTING
 from psutil.tests import COVERAGE
 from psutil.tests import HAS_CONNECTIONS_UNIX
 from psutil.tests import PYTHON_EXE
+from psutil.tests import PYTHON_EXE_ENV
 from psutil.tests import PsutilTestCase
 from psutil.tests import TestMemoryLeak
 from psutil.tests import bind_socket
@@ -260,7 +261,8 @@ class TestProcessUtils(PsutilTestCase):
         terminate(p)
         # by psutil.Popen
         cmd = [PYTHON_EXE, "-c", "import time; time.sleep(60);"]
-        p = psutil.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
+        p = psutil.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
+                         env=PYTHON_EXE_ENV)
         terminate(p)
         self.assertProcessGone(p)
         terminate(p)