Fix virtualenv environment detection

There seems to be some combination of factors that cause the string
matching of the VIRTUAL_ENV envvar to fail on trybots/pre-cq.  Given
that, I think its generally a bad idea to do path comparisons against
the virtualenv environment and reverting back to a dumber, but more
robust implementation of the virtualenv check using a magic variable.

BUG=chromium:673926
TEST=Run bin/venv_check and unit tests

Change-Id: Ifaafafc0abf5b079585a168f5d6a4b731c1a2d2e
Reviewed-on: https://chromium-review.googlesource.com/419722
Commit-Ready: David Riley <davidriley@chromium.org>
Tested-by: Allen Li <ayatane@chromium.org>
Reviewed-by: Aviv Keshet <akeshet@chromium.org>
diff --git a/scripts/virtualenv_wrapper.py b/scripts/virtualenv_wrapper.py
index 61bb7b0..15def65 100755
--- a/scripts/virtualenv_wrapper.py
+++ b/scripts/virtualenv_wrapper.py
@@ -13,6 +13,7 @@
 
 import wrapper
 
+
 def _FindChromiteDir():
   path = os.path.dirname(os.path.realpath(__file__))
   while not os.path.exists(os.path.join(path, 'PRESUBMIT.cfg')):
@@ -21,15 +22,21 @@
 
 
 _CHROMITE_DIR = _FindChromiteDir()
+
 # _VIRTUALENV_DIR contains the scripts for working with venvs
 _VIRTUALENV_DIR = os.path.join(_CHROMITE_DIR, '../infra_virtualenv')
-# _VENV_DIR is the actual virtualenv that contains bin/activate.
+_CREATE_VENV_PATH = os.path.join(_VIRTUALENV_DIR, 'create_venv')
+_VENV_COMMAND_PATH = os.path.join(_VIRTUALENV_DIR, 'venv_command')
+
+# _VENV_DIR is the virtualenv dir that contains bin/activate.
 _VENV_DIR = os.path.join(_CHROMITE_DIR, '.venv')
 _REQUIREMENTS = os.path.join(_CHROMITE_DIR, 'venv', 'requirements.txt')
 
+_VENV_MARKER = 'INSIDE_CHROMITE_VENV'
+
 
 def main():
-  if _IsInsideVenv():
+  if _IsInsideVenv(os.environ):
     wrapper.DoMain()
   else:
     _CreateVenv()
@@ -39,7 +46,7 @@
 def _CreateVenv():
   """Create or update chromite venv."""
   subprocess.check_call([
-      os.path.join(_VIRTUALENV_DIR, 'create_venv'),
+      _CREATE_VENV_PATH,
       _VENV_DIR,
       _REQUIREMENTS,
   ], stdout=sys.stderr)
@@ -51,15 +58,43 @@
   Args:
     args: Sequence of arguments.
   """
-  os.execv(os.path.join(_VIRTUALENV_DIR, 'venv_command'),
-           ['venv_command', _VENV_DIR] + list(args))
+  os.execve(_VENV_COMMAND_PATH,
+            [os.path.basename(_VENV_COMMAND_PATH), _VENV_DIR] + list(args),
+            _CreateVenvEnvironment(os.environ))
 
 
-def _IsInsideVenv():
-  """Return whether we're running inside a virtualenv."""
-  # Proper way is checking sys.prefix and sys.base_prefix in Python 3.
-  # PEP 405 isn't fully implemented in Python 2.
-  return _VENV_DIR in os.environ.get('VIRTUAL_ENV', '')
+def _CreateVenvEnvironment(env_dict):
+  """Create environment for a virtualenv.
+
+  This adds a special marker variable to a copy of the input environment dict
+  and returns the copy.
+
+  Args:
+    env_dict: Environment variable dict to use as base, which is not modified.
+
+  Returns:
+    New environment dict for a virtualenv.
+  """
+  new_env_dict = env_dict.copy()
+  new_env_dict[_VENV_MARKER] = '1'
+  return new_env_dict
+
+
+def _IsInsideVenv(env_dict):
+  """Return whether the environment dict is running inside a virtualenv.
+
+  This checks the environment dict for the special marker added by
+  _CreateVenvEnvironment().
+
+  Args:
+    env_dict: Environment variable dict to check
+
+  Returns:
+    A true value if inside virtualenv, else a false value.
+  """
+  # Checking sys.prefix or doing any kind of path check is unreliable because
+  # we check out chromite to weird places.
+  return env_dict.get(_VENV_MARKER, '')
 
 
 if __name__ == '__main__':