Cleanup run-test build environment passing.

Use python variables rather than environment to pass state.

Cleanup the environment use left over from the bash version.

Test: build artifacts are same as before
Change-Id: I3ba7f9e31a9a7046c62d67c6c66af4988d5c469a
diff --git a/test/004-JniTest/build.py b/test/004-JniTest/build.py
index 26745ab..4bb504d 100644
--- a/test/004-JniTest/build.py
+++ b/test/004-JniTest/build.py
@@ -33,6 +33,6 @@
 
   # Remove the *-aotex build artifacts (but keep src-aotex) with dalvik.* annotations.
   shutil.rmtree("classes-aotex")
-  if os.environ["BUILD_MODE"] != "jvm":
+  if not ctx.jvm:
     os.remove("classes-aotex.jar")
     os.remove("004-JniTest-aotex.jar")
diff --git a/test/089-many-methods/build.py b/test/089-many-methods/build.py
index 9710f32..414d6de 100644
--- a/test/089-many-methods/build.py
+++ b/test/089-many-methods/build.py
@@ -21,7 +21,8 @@
   # Force DEX generation so test also passes with --jvm.
   try:
     ctx.default_build(api_level=20, need_dex=True)
-    assert False, "Test was not expected to build successfully"
   except Exception as e:
     # Check that a build failure happened (the test is not expected to run).
     assert "Cannot fit requested classes in a single dex" in str(e), e
+    return
+  assert False, "Test was not expected to build successfully"
diff --git a/test/166-bad-interface-super/build.py b/test/166-bad-interface-super/build.py
index e4502b3..ce7bd34 100644
--- a/test/166-bad-interface-super/build.py
+++ b/test/166-bad-interface-super/build.py
@@ -18,7 +18,7 @@
 
 def build(ctx):
   # Use the jasmin sources for JVM, otherwise the smali sources.
-  if os.environ["BUILD_MODE"] == "jvm":
+  if ctx.jvm:
     ctx.default_build(has_smali=False)
   else:
     ctx.default_build(has_jasmin=False)
diff --git a/test/180-native-default-method/build.py b/test/180-native-default-method/build.py
index 46c3c5c..5997ee5 100644
--- a/test/180-native-default-method/build.py
+++ b/test/180-native-default-method/build.py
@@ -18,16 +18,15 @@
 
 def build(ctx):
   ctx.default_build()
-
-  if os.environ["BUILD_MODE"] != "jvm":
-    # Change the generated dex file to have a v35 magic number if it is version 38
-    with open("classes.dex", "rb+") as f:
-      assert f.read(8) == b"dex\n038\x00"
-      f.seek(0)
-      f.write(b"dex\n035\x00")
-    os.remove("180-native-default-method.jar")
-    cmd = [
-        os.environ["SOONG_ZIP"], "-o", "180-native-default-method.jar", "-f",
-        "classes.dex"
-    ]
-    subprocess.run(cmd, check=True)
+  if ctx.jvm:
+    return
+  # Change the generated dex file to have a v35 magic number if it is version 38
+  with open("classes.dex", "rb+") as f:
+    assert f.read(8) == b"dex\n038\x00"
+    f.seek(0)
+    f.write(b"dex\n035\x00")
+  os.remove("180-native-default-method.jar")
+  cmd = [
+      ctx.soong_zip, "-o", "180-native-default-method.jar", "-f", "classes.dex"
+  ]
+  subprocess.run(cmd, check=True)
diff --git a/test/1965-get-set-local-primitive-no-tables/build.py b/test/1965-get-set-local-primitive-no-tables/build.py
index b7a608b..3a908c0 100644
--- a/test/1965-get-set-local-primitive-no-tables/build.py
+++ b/test/1965-get-set-local-primitive-no-tables/build.py
@@ -17,5 +17,5 @@
 
 
 def build(ctx):
-  ctx.bash("./generate-sources --" + os.environ["BUILD_MODE"])
+  ctx.bash("./generate-sources --" + ctx.mode)
   ctx.default_build()
diff --git a/test/1966-get-set-local-objects-no-table/build.py b/test/1966-get-set-local-objects-no-table/build.py
index b7a608b..3a908c0 100644
--- a/test/1966-get-set-local-objects-no-table/build.py
+++ b/test/1966-get-set-local-objects-no-table/build.py
@@ -17,5 +17,5 @@
 
 
 def build(ctx):
-  ctx.bash("./generate-sources --" + os.environ["BUILD_MODE"])
+  ctx.bash("./generate-sources --" + ctx.mode)
   ctx.default_build()
diff --git a/test/370-dex-v37/build.py b/test/370-dex-v37/build.py
index 1fa6912..b9e3ffa 100644
--- a/test/370-dex-v37/build.py
+++ b/test/370-dex-v37/build.py
@@ -18,16 +18,13 @@
 
 def build(ctx):
   ctx.default_build()
-
-  if os.environ["BUILD_MODE"] != "jvm":
-    # Change the generated dex file to have a v37 magic number if it is version 35
-    with open("classes.dex", "rb+") as f:
-      if f.read(8) == b"dex\n035\x00":
-        f.seek(0)
-        f.write(b"dex\n037\x00")
-        os.remove("370-dex-v37.jar")
-    cmd = [
-        os.environ["SOONG_ZIP"], "-o", "370-dex-v37.jar", "-f",
-        "classes.dex"
-    ]
-    subprocess.run(cmd, check=True)
+  if ctx.jvm:
+    return
+  # Change the generated dex file to have a v37 magic number if it is version 35
+  with open("classes.dex", "rb+") as f:
+    if f.read(8) == b"dex\n035\x00":
+      f.seek(0)
+      f.write(b"dex\n037\x00")
+      os.remove("370-dex-v37.jar")
+  cmd = [ctx.soong_zip, "-o", "370-dex-v37.jar", "-f", "classes.dex"]
+  subprocess.run(cmd, check=True)
diff --git a/test/968-default-partial-compile-gen/build.py b/test/968-default-partial-compile-gen/build.py
index 8cbdf95..8868b70 100644
--- a/test/968-default-partial-compile-gen/build.py
+++ b/test/968-default-partial-compile-gen/build.py
@@ -17,6 +17,7 @@
 
 
 def build(ctx):
-  ctx.bash("./generate-sources --" + os.environ["BUILD_MODE"])
-  if os.environ["BUILD_MODE"] != "jvm":
-    ctx.default_build(experimental="default-methods")
+  ctx.bash("./generate-sources --" + ctx.mode)
+  if ctx.jvm:
+    return
+  ctx.default_build(experimental="default-methods")
diff --git a/test/970-iface-super-resolution-gen/build.py b/test/970-iface-super-resolution-gen/build.py
index 57abccd..795eb6e 100644
--- a/test/970-iface-super-resolution-gen/build.py
+++ b/test/970-iface-super-resolution-gen/build.py
@@ -17,5 +17,5 @@
 
 
 def build(ctx):
-  ctx.bash("./generate-sources --" + os.environ["BUILD_MODE"])
+  ctx.bash("./generate-sources --" + ctx.mode)
   ctx.default_build(experimental="default-methods")
diff --git a/test/971-iface-super/build.py b/test/971-iface-super/build.py
index 8cbdf95..8868b70 100644
--- a/test/971-iface-super/build.py
+++ b/test/971-iface-super/build.py
@@ -17,6 +17,7 @@
 
 
 def build(ctx):
-  ctx.bash("./generate-sources --" + os.environ["BUILD_MODE"])
-  if os.environ["BUILD_MODE"] != "jvm":
-    ctx.default_build(experimental="default-methods")
+  ctx.bash("./generate-sources --" + ctx.mode)
+  if ctx.jvm:
+    return
+  ctx.default_build(experimental="default-methods")
diff --git a/test/art_build_rules.py b/test/art_build_rules.py
index b1e81be..bd954d2 100644
--- a/test/art_build_rules.py
+++ b/test/art_build_rules.py
@@ -32,10 +32,60 @@
 from shutil import rmtree
 from os import remove
 from re import match
+from os.path import join
 
 USE_RBE_FOR_JAVAC = 100    # Percentage of tests that can use RBE (between 0 and 100)
 USE_RBE_FOR_D8 = 100       # Percentage of tests that can use RBE (between 0 and 100)
 
+class BuildTestContext:
+  def __init__(self, args, build_top, sbox, test_name, test_dir):
+    self.test_dir = test_dir
+    self.mode = args.mode
+    self.jvm = (self.mode == "jvm")
+    self.host = (self.mode == "host")
+    self.target = (self.mode == "target")
+    assert self.jvm or self.host or self.target
+
+    java_home = os.environ.get("JAVA_HOME")
+    tools_dir = os.path.abspath(join(os.path.dirname(__file__), "../../../out/bin"))
+    self.android_build_top = build_top
+    self.art_test_run_test_bootclasspath = join(build_top, args.bootclasspath)
+    self.d8 = join(tools_dir, "d8")
+    self.d8_flags = []
+    self.hiddenapi = join(tools_dir, "hiddenapi")
+    self.jasmin = join(tools_dir, "jasmin")
+    self.java = join(java_home, "bin/java")
+    self.javac = join(java_home, "bin/javac")
+    self.javac_args = "-g -Xlint:-options -source 1.8 -target 1.8"
+    self.need_dex = (self.host or self.target)
+    self.sbox_path = sbox
+    self.smali = join(tools_dir, "smali")
+    self.smali_flags = []
+    self.soong_zip = join(build_top, "prebuilts/build-tools/linux-x86/bin/soong_zip")
+    self.test_name = test_name
+    self.zipalign = join(build_top, "prebuilts/build-tools/linux-x86/bin/zipalign")
+
+    # Minimal environment needed for bash commands that we execute.
+    self.bash_env = {
+      "ANDROID_BUILD_TOP": self.android_build_top,
+      "D8": self.d8,
+      "JAVA": self.java,
+      "JAVAC": self.javac,
+      "JAVAC_ARGS": self.javac_args,
+      "JAVA_HOME": java_home,
+      "PATH": os.environ["PATH"],
+      "PYTHONDONTWRITEBYTECODE": "1",
+      "SMALI": self.smali,
+      "SOONG_ZIP": self.soong_zip,
+      "TEST_NAME": self.test_name,
+    }
+
+  def bash(self, cmd):
+    return subprocess.run(cmd, shell=True, env=self.bash_env, check=True)
+
+  def default_build(self, **kwargs):
+    globals()['default_build'](self, **kwargs)
+
 def rm(*patterns):
   for pattern in patterns:
     for path in glob.glob(pattern):
@@ -60,15 +110,12 @@
     has_jasmin=None,
   ):
 
-  def parse_bool(text):
-    return {"true": True, "false": False}[text.lower()]
-
-  ANDROID_BUILD_TOP = os.environ["ANDROID_BUILD_TOP"]
-  SBOX_PATH = os.environ["SBOX_PATH"]
+  ANDROID_BUILD_TOP = ctx.android_build_top
+  SBOX_PATH = ctx.sbox_path
   CWD = os.getcwd()
-  TEST_NAME = os.environ["TEST_NAME"]
-  ART_TEST_RUN_TEST_BOOTCLASSPATH = path.relpath(os.environ["ART_TEST_RUN_TEST_BOOTCLASSPATH"], CWD)
-  NEED_DEX = parse_bool(os.environ["NEED_DEX"]) if need_dex is None else need_dex
+  TEST_NAME = ctx.test_name
+  ART_TEST_RUN_TEST_BOOTCLASSPATH = path.relpath(ctx.art_test_run_test_bootclasspath, CWD)
+  NEED_DEX = ctx.need_dex if need_dex is None else need_dex
 
   RBE_exec_root = os.environ.get("RBE_exec_root")
   RBE_rewrapper = path.join(ANDROID_BUILD_TOP, "prebuilts/remoteexecution-client/live/rewrapper")
@@ -89,11 +136,11 @@
   HAS_SRC_BCPEX = path.exists("src-bcpex")
   HAS_HIDDENAPI_SPEC = path.exists("hiddenapi-flags.csv")
 
-  JAVAC_ARGS = shlex.split(os.environ.get("JAVAC_ARGS", "")) + javac_args
-  SMALI_ARGS = shlex.split(os.environ.get("SMALI_ARGS", "")) + smali_args
-  D8_FLAGS = shlex.split(os.environ.get("D8_FLAGS", "")) + d8_flags
+  JAVAC_ARGS = shlex.split(ctx.javac_args) + javac_args
+  SMALI_ARGS = ctx.smali_flags + smali_args
+  D8_FLAGS = ctx.d8_flags + d8_flags
 
-  BUILD_MODE = os.environ["BUILD_MODE"]
+  BUILD_MODE = ctx.mode
 
   # Setup experimental API level mappings in a bash associative array.
   EXPERIMENTAL_API_LEVEL = {}
@@ -119,13 +166,15 @@
   SMALI_ARGS.extend(["--api", str(api_level)])
   D8_FLAGS.extend(["--min-api", str(api_level)])
 
-
   def run(executable, args):
     cmd = shlex.split(executable) + args
     if executable.endswith(".sh"):
       cmd = ["/bin/bash"] + cmd
+    env = ctx.bash_env
+    env.update({k: v for k, v in os.environ.items() if k.startswith("RBE_")})
     p = subprocess.run(cmd,
                        encoding=os.sys.stdout.encoding,
+                       env=ctx.bash_env,
                        stderr=subprocess.STDOUT,
                        stdout=subprocess.PIPE)
     if p.returncode != 0:
@@ -135,13 +184,13 @@
 
 
   # Helper functions to execute tools.
-  soong_zip = functools.partial(run, os.environ["SOONG_ZIP"])
-  zipalign = functools.partial(run, os.environ["ZIPALIGN"])
-  javac = functools.partial(run, os.environ["JAVAC"])
-  jasmin = functools.partial(run, os.environ["JASMIN"])
-  smali = functools.partial(run, os.environ["SMALI"])
-  d8 = functools.partial(run, os.environ["D8"])
-  hiddenapi = functools.partial(run, os.environ["HIDDENAPI"])
+  soong_zip = functools.partial(run, ctx.soong_zip)
+  zipalign = functools.partial(run, ctx.zipalign)
+  javac = functools.partial(run, ctx.javac)
+  jasmin = functools.partial(run, ctx.jasmin)
+  smali = functools.partial(run, ctx.smali)
+  d8 = functools.partial(run, ctx.d8)
+  hiddenapi = functools.partial(run, ctx.hiddenapi)
 
   if "RBE_server_address" in os.environ:
     version = match(r"Version: (\d*)\.(\d*)\.(\d*)", run(RBE_rewrapper, ["--version"]).stdout)
@@ -164,7 +213,7 @@
         output = path.relpath(path.join(CWD, args[args.index("-d") + 1]), RBE_exec_root)
         return rbe_wrap([
           "--output_directories", output,
-          os.path.relpath(os.environ["JAVAC"], CWD),
+          os.path.relpath(ctx.javac, CWD),
         ] + args)
 
     if USE_RBE_FOR_D8 > (hash(TEST_NAME) % 100):  # Use for given percentage of tests.
@@ -174,7 +223,7 @@
         return rbe_wrap([
           "--output_files" if output.endswith(".jar") else "--output_directories", output,
           "--toolchain_inputs=prebuilts/jdk/jdk11/linux-x86/bin/java",
-          os.path.relpath(os.environ["D8"], CWD)] + args, inputs)
+          os.path.relpath(ctx.d8, CWD)] + args, inputs)
 
   # If wrapper script exists, use it instead of the default javac.
   if os.path.exists("javac_wrapper.sh"):
diff --git a/test/run-test-build.py b/test/run-test-build.py
index 4074015..30a966d 100755
--- a/test/run-test-build.py
+++ b/test/run-test-build.py
@@ -19,113 +19,96 @@
 It is intended to be used only from soong genrule.
 """
 
-import argparse, os, shutil, subprocess, glob, re, json, multiprocessing, pathlib, fcntl
-import art_build_rules
+from argparse import ArgumentParser
+from art_build_rules import BuildTestContext, default_build
+from fcntl import lockf, LOCK_EX, LOCK_NB
 from importlib.machinery import SourceFileLoader
+from multiprocessing import Pool
+from multiprocessing.pool import ApplyResult
+from os import environ, getcwd, chdir, cpu_count
 from os.path import join, basename
-
-import art_build_rules
+from pathlib import Path
+from re import match
+from shutil import copytree
+from subprocess import run
+from typing import Dict
 
 ZIP = "prebuilts/build-tools/linux-x86/bin/soong_zip"
 
-class BuildTestContext:
-  def __init__(self, mode):
-    self.jvm = (mode == "jvm")
-    self.host = (mode == "host")
-    self.target = (mode == "target")
+lock_file = None  # Keep alive as long as this process is alive.
 
-  def bash(self, cmd):
-    return subprocess.run(cmd, shell=True, check=True)
 
-  def default_build(self, **kwargs):
-    art_build_rules.default_build(self, **kwargs)
-
-def copy_sources(args, tmp, mode, srcdir):
+def copy_sources(args, ziproot: Path, mode: str, srcdir: Path) -> Path:
   """Copy test files from Android tree into the build sandbox and return its path."""
 
-  dstdir = join(tmp, mode, basename(srcdir))
-  shutil.copytree(srcdir, dstdir)
+  dstdir = ziproot / mode / srcdir.name
+  copytree(srcdir, dstdir)
   return dstdir
 
-def build_test(args, mode, build_top, sbox, dstdir):
+
+def build_test(ctx: BuildTestContext) -> None:
   """Run the build script for single run-test"""
 
-  join = os.path.join
-  java_home = os.environ.get("JAVA_HOME")
-  tools_dir = os.path.abspath(join(os.path.dirname(__file__), "../../../out/bin"))
-  test_name = os.path.basename(dstdir)
-  env = dict(os.environ)
-  env.update({
-    "BUILD_MODE": mode,
-    "ANDROID_BUILD_TOP": build_top,
-    "SBOX_PATH":   sbox,
-    "ART_TEST_RUN_TEST_BOOTCLASSPATH": join(build_top, args.bootclasspath),
-    "TEST_NAME":   test_name,
-    "SOONG_ZIP":   join(build_top, "prebuilts/build-tools/linux-x86/bin/soong_zip"),
-    "ZIPALIGN":    join(build_top, "prebuilts/build-tools/linux-x86/bin/zipalign"),
-    "JAVA":        join(java_home, "bin/java"),
-    "JAVAC":       join(java_home, "bin/javac"),
-    "JAVAC_ARGS":  "-g -Xlint:-options -source 1.8 -target 1.8",
-    "D8":          join(tools_dir, "d8"),
-    "HIDDENAPI":   join(tools_dir, "hiddenapi"),
-    "JASMIN":      join(tools_dir, "jasmin"),
-    "SMALI":       join(tools_dir, "smali"),
-    "NEED_DEX":    {"host": "true", "target": "true", "jvm": "false"}[mode],
-  })
-
-  os.chdir(dstdir)
-  for name, value in env.items():
-    os.environ[name] = str(value)
-  ctx = BuildTestContext(mode)
-  script = pathlib.Path(join(dstdir, "build.py"))
+  chdir(ctx.test_dir)
+  script = ctx.test_dir / "build.py"
   if script.exists():
-    module = SourceFileLoader("build_" + test_name, str(script)).load_module()
+    module = SourceFileLoader("build_" + ctx.test_name,
+                              str(script)).load_module()
     module.build(ctx)
   else:
-    art_build_rules.default_build(ctx)
+    default_build(ctx)
+
 
 # If we build just individual shard, we want to split the work among all the cores,
 # but if the build system builds all shards, we don't want to overload the machine.
 # We don't know which situation we are in, so as simple work-around, we use a lock
 # file to allow only one shard to use multiprocessing at the same time.
-def use_multiprocessing(mode):
-  global lock_file  # Keep alive as long as this process is alive.
-  lock_path = os.path.join(os.environ["TMPDIR"], "art-test-run-test-build-py-" + mode)
+def use_multiprocessing(mode: str) -> bool:
+  global lock_file
+  lock_path = join(environ["TMPDIR"], "art-test-run-test-build-py-" + mode)
   lock_file = open(lock_path, "w")
   try:
-    fcntl.lockf(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
+    lockf(lock_file, LOCK_EX | LOCK_NB)
     return True  # We are the only instance of this script in the build system.
   except BlockingIOError:
     return False  # Some other instance is already running.
 
-def main():
-  parser = argparse.ArgumentParser(description=__doc__)
-  parser.add_argument("--out", help="Path of the generated ZIP file with the build data")
-  parser.add_argument('--mode', choices=['host', 'jvm', 'target'])
-  parser.add_argument("--shard", help="Identifies subset of tests to build (00..99)")
-  parser.add_argument("--bootclasspath", help="JAR files used for javac compilation")
+
+def main() -> None:
+  parser = ArgumentParser(description=__doc__)
+  parser.add_argument(
+      "--out", help="Path of the generated ZIP file with the build data")
+  parser.add_argument("--mode", choices=["host", "jvm", "target"])
+  parser.add_argument(
+      "--shard", help="Identifies subset of tests to build (00..99)")
+  parser.add_argument(
+      "--bootclasspath", help="JAR files used for javac compilation")
   args = parser.parse_args()
 
-  build_top = os.getcwd()
-  sbox = pathlib.Path(__file__).absolute().parent.parent.parent.parent.parent
+  build_top = Path(getcwd())
+  sbox = Path(__file__).absolute().parent.parent.parent.parent.parent
   assert sbox.parent.name == "sbox" and len(sbox.name) == 40
 
-  ziproot = os.path.join(sbox, "zip")
-  srcdirs = sorted(glob.glob(os.path.join("art", "test", "*")))
-  srcdirs = filter(lambda srcdir: re.match(".*/\d*{}-.*".format(args.shard), srcdir), srcdirs)
-  dstdirs = [copy_sources(args, ziproot, args.mode, srcdir) for srcdir in srcdirs]
-  dstdirs = filter(lambda dstdir: dstdir, dstdirs)  # Remove None (skipped tests).
-  # Use multiprocess (i.e. forking) since tests modify their current working directory.
-  with multiprocessing.Pool(os.cpu_count() if use_multiprocessing(args.mode) else 1) as pool:
-    jobs = [(d, pool.apply_async(build_test, (args, args.mode, build_top, sbox, d))) for d in dstdirs]
-    for dstdir, job in jobs:
+  ziproot = sbox / "zip"
+  srcdirs = sorted(build_top.glob("art/test/*"))
+  srcdirs = [s for s in srcdirs if match("\d*{}-.*".format(args.shard), s.name)]
+  dstdirs = [copy_sources(args, ziproot, args.mode, s) for s in srcdirs]
+
+  # Use multiprocessing (i.e. forking) since tests modify their current working directory.
+  with Pool(cpu_count() if use_multiprocessing(args.mode) else 1) as pool:
+    jobs: Dict[Path, ApplyResult] = {}
+    for dstdir in dstdirs:
+      ctx = BuildTestContext(args, build_top, sbox, dstdir.name, dstdir)
+      jobs[dstdir] = pool.apply_async(build_test, (ctx,))
+    for dstdir, job in jobs.items():
       try:
         job.get()
       except Exception as e:
-        raise Exception("Failed to build " + os.path.basename(dstdir)) from e.__cause__
+        raise Exception("Failed to build " + dstdir.name) from e.__cause__
 
   # Create the final zip file which contains the content of the temporary directory.
-  proc = subprocess.run([ZIP, "-o", args.out, "-C", ziproot, "-D", ziproot], check=True)
+  proc = run([ZIP, "-o", args.out, "-C", ziproot, "-D", ziproot], check=True)
+
 
 if __name__ == "__main__":
   main()