From 3867c4f95f6b78e7d2232d57ed91891c390b66e4 Mon Sep 17 00:00:00 2001
From: Rob Mohr <mohrr@google.com>
Date: Mon, 2 Mar 2020 09:46:20 -0800
Subject: [PATCH] Add check for virtual env to doctor

Change-Id: If2aca9af65809828a8070603c618100856c6580f
---
 pw_doctor/py/pw_doctor/doctor.py | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/pw_doctor/py/pw_doctor/doctor.py b/pw_doctor/py/pw_doctor/doctor.py
index d45fe0581..ae0ff385e 100755
--- a/pw_doctor/py/pw_doctor/doctor.py
+++ b/pw_doctor/py/pw_doctor/doctor.py
@@ -18,6 +18,7 @@ import argparse
 import logging
 import json
 import os
+import pathlib
 import shutil
 import subprocess
 import sys
@@ -104,6 +105,27 @@ def python_version(ctx: DoctorContext):
                     *actual[0:3])
 
 
+@register_into(CHECKS)
+def virtualenv(ctx: DoctorContext):
+    """Check that we're in the correct virtualenv."""
+    try:
+        venv_path = pathlib.Path(os.environ['VIRTUAL_ENV']).resolve()
+    except KeyError:
+        ctx.error('VIRTUAL_ENV not set')
+        return
+
+    # When running in LUCI we might not have gone through the normal environment
+    # setup process, so we need to skip the rest of this step.
+    if 'LUCI_CONTEXT' in os.environ:
+        return
+
+    root = pathlib.Path(os.environ['PW_ROOT']).resolve()
+
+    if root not in venv_path.parents:
+        ctx.error('VIRTUAL_ENV (%s) not inside PW_ROOT (%s)', venv_path, root)
+        ctx.error('\n'.join(os.environ.keys()))
+
+
 @register_into(CHECKS)
 def cipd(ctx: DoctorContext):
     """Check cipd is set up correctly and in use."""