compiler_wrapper: add python search path

Exporting BISECT_STAGE to POPULATE_GOOD currently makes the wrapper fail
to build any files with the error message of "ModuleNotFoundError: No
module named 'bisect_driver'". This patch fixes the issue by adding
search path for bisect_driver.py explicitly.

BUG=chromium:1042452

TEST=verified locally.

Change-Id: I7c73459b377d9011c7087a9e241db378ffb216e8
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/2003947
Tested-by: Jian Cai <jiancai@google.com>
Reviewed-by: Manoj Gupta <manojgupta@chromium.org>
Reviewed-by: George Burgess <gbiv@chromium.org>

8 files changed, 63 insertions, 10 deletions

diff --git a/compiler_wrapper/bisect_flag.go b/compiler_wrapper/bisect_flag.go
index 6271e23f..8aec3e7e 100644
--- a/compiler_wrapper/bisect_flag.go
+++ b/compiler_wrapper/bisect_flag.go
@@ -6,6 +6,7 @@ package main
 
 import (
 	"errors"
+	"os"
 	"path/filepath"
 )
 
@@ -51,6 +52,15 @@ func calcBisectCommand(env env, cfg *config, bisectStage string, compilerCmd *co
 		}
 	}
 	absCompilerPath := getAbsCmdPath(env, compilerCmd)
+	pythonPath, err := filepath.Abs(os.Args[0])
+	if err != nil {
+		return nil, err
+	}
+	pythonPath, err = filepath.EvalSymlinks(pythonPath)
+	if err != nil {
+		return nil, err
+	}
+	pythonPath = filepath.Dir(pythonPath)
 	return &command{
 		Path: "/usr/bin/env",
 		Args: append([]string{
@@ -61,6 +71,6 @@ func calcBisectCommand(env env, cfg *config, bisectStage string, compilerCmd *co
 			bisectDir,
 			absCompilerPath,
 		}, compilerCmd.Args...),
-		EnvUpdates: compilerCmd.EnvUpdates,
+		EnvUpdates: append(compilerCmd.EnvUpdates, "PYTHONPATH="+pythonPath),
 	}, nil
 }
diff --git a/compiler_wrapper/goldenutil_test.go b/compiler_wrapper/goldenutil_test.go
index e6b42f77..23e003c5 100644
--- a/compiler_wrapper/goldenutil_test.go
+++ b/compiler_wrapper/goldenutil_test.go
@@ -14,6 +14,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"runtime"
 	"strings"
 )
 
@@ -145,7 +146,13 @@ func fillGoldenResults(ctx *testContext, files []goldenFile) []goldenFile {
 			// Create an empty wrapper at the given path.
 			// Needed as we are resolving symlinks which stats the wrapper file.
 			ctx.writeFile(record.WrapperCmd.Cmd.Path, "")
+			// Assign a fixed path to os.Args[0] to pass the test.
+			tmp := os.Args[0]
+			// callCompiler verifies os.Args[0] exists, so use a real file.
+			_, file, _, _ := runtime.Caller(1)
+			os.Args[0] = file
 			record.WrapperCmd.ExitCode = callCompiler(ctx, ctx.cfg, record.WrapperCmd.Cmd)
+			os.Args[0] = tmp
 			if hasInternalError(ctx.stderrString()) {
 				ctx.t.Errorf("found an internal error for wrapperCmd %#v and env #%v. Got: %s",
 					record.WrapperCmd.Cmd, record.Env, ctx.stderrString())
diff --git a/compiler_wrapper/testdata/android_golden/bisect.json b/compiler_wrapper/testdata/android_golden/bisect.json
index a24222ab..be07da68 100644
--- a/compiler_wrapper/testdata/android_golden/bisect.json
+++ b/compiler_wrapper/testdata/android_golden/bisect.json
@@ -25,6 +25,9 @@
             "/user/home/ANDROID_BISECT",
             "/tmp/stable/clang.real",
             "main.cc"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -57,6 +60,9 @@
             "someBisectDir",
             "/tmp/stable/clang.real",
             "main.cc"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -92,6 +98,9 @@
             "someBisectDir",
             "/tmp/stable/clang.real",
             "main.cc"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         },
         "stdout": "somemessage",
diff --git a/compiler_wrapper/testdata/cros_clang_host_golden/bisect.json b/compiler_wrapper/testdata/cros_clang_host_golden/bisect.json
index 5920a365..ad4feac3 100644
--- a/compiler_wrapper/testdata/cros_clang_host_golden/bisect.json
+++ b/compiler_wrapper/testdata/cros_clang_host_golden/bisect.json
@@ -39,6 +39,9 @@
             "-Wno-unknown-warning-option",
             "main.cc",
             "-Wno-implicit-int-float-conversion"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -85,6 +88,9 @@
             "-Wno-unknown-warning-option",
             "main.cc",
             "-Wno-implicit-int-float-conversion"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -134,6 +140,9 @@
             "-Wno-unknown-warning-option",
             "main.cc",
             "-Wno-implicit-int-float-conversion"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         },
         "stdout": "somemessage",
diff --git a/compiler_wrapper/testdata/cros_hardened_golden/bisect.json b/compiler_wrapper/testdata/cros_hardened_golden/bisect.json
index c1daaa88..6e1d4981 100644
--- a/compiler_wrapper/testdata/cros_hardened_golden/bisect.json
+++ b/compiler_wrapper/testdata/cros_hardened_golden/bisect.json
@@ -54,7 +54,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -116,7 +117,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -181,7 +183,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         },
         "stdout": "somemessage",
diff --git a/compiler_wrapper/testdata/cros_hardened_llvmnext_golden/bisect.json b/compiler_wrapper/testdata/cros_hardened_llvmnext_golden/bisect.json
index c1daaa88..6e1d4981 100644
--- a/compiler_wrapper/testdata/cros_hardened_llvmnext_golden/bisect.json
+++ b/compiler_wrapper/testdata/cros_hardened_llvmnext_golden/bisect.json
@@ -54,7 +54,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -116,7 +117,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -181,7 +183,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         },
         "stdout": "somemessage",
diff --git a/compiler_wrapper/testdata/cros_hardened_noccache_golden/bisect.json b/compiler_wrapper/testdata/cros_hardened_noccache_golden/bisect.json
index 46c29635..3fafcf01 100644
--- a/compiler_wrapper/testdata/cros_hardened_noccache_golden/bisect.json
+++ b/compiler_wrapper/testdata/cros_hardened_noccache_golden/bisect.json
@@ -48,6 +48,9 @@
             "-B../../bin",
             "-target",
             "x86_64-cros-linux-gnu"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -103,6 +106,9 @@
             "-B../../bin",
             "-target",
             "x86_64-cros-linux-gnu"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -161,6 +167,9 @@
             "-B../../bin",
             "-target",
             "x86_64-cros-linux-gnu"
+          ],
+          "env_updates": [
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         },
         "stdout": "somemessage",
diff --git a/compiler_wrapper/testdata/cros_nonhardened_golden/bisect.json b/compiler_wrapper/testdata/cros_nonhardened_golden/bisect.json
index b97c379d..7d2acdac 100644
--- a/compiler_wrapper/testdata/cros_nonhardened_golden/bisect.json
+++ b/compiler_wrapper/testdata/cros_nonhardened_golden/bisect.json
@@ -46,7 +46,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -100,7 +101,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         }
       }
@@ -157,7 +159,8 @@
             "CCACHE_BASEDIR=/usr/x86_64-cros-linux-gnu",
             "CCACHE_DIR=/var/cache/distfiles/ccache",
             "CCACHE_UMASK=002",
-            "CCACHE_CPP2=yes"
+            "CCACHE_CPP2=yes",
+            "PYTHONPATH=/media/storage/jiancai/chromeos/src/third_party/toolchain-utils/compiler_wrapper"
           ]
         },
         "stdout": "somemessage",