run_pytest: Add namespacing & chroot initialization

Our tests need a chroot and want to be run inside a namespace in order
to not clobber other resources. The previous test runner was responsible
for setting all of this up, so port this behavior over to the new test
wrapper.

BUG=chromium:1064762, chromium:1064760, chromium:1064761
TEST=`run_pytest` inside and outside chroot

Change-Id: If85a93eaccf13171babb2ba797dfe93982474580
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/chromite/+/2130751
Tested-by: Chris McDonald <cjmcdonald@chromium.org>
Commit-Queue: Chris McDonald <cjmcdonald@chromium.org>
Reviewed-by: Mike Frysinger <vapier@chromium.org>
diff --git a/scripts/run_pytest.py b/scripts/run_pytest.py
index 8772ff7..5655368 100644
--- a/scripts/run_pytest.py
+++ b/scripts/run_pytest.py
@@ -7,10 +7,69 @@
 
 from __future__ import print_function
 
+import os
 import sys
 
 import pytest  # pylint: disable=import-error
 
+from chromite.lib import cgroups
+from chromite.lib import constants
+from chromite.lib import cros_build_lib
+from chromite.lib import gs
+from chromite.lib import namespaces
+
 
 def main(argv):
+  ensure_chroot_exists()
+  re_execute_inside_chroot(argv)
+
+  # This is a cheesy hack to make sure gsutil is populated in the cache before
+  # we run tests. This is a partial workaround for crbug.com/468838.
+  gs.GSContext.GetDefaultGSUtilBin()
+
+  re_execute_with_namespace([sys.argv[0]] + argv)
+
   sys.exit(pytest.main(argv))
+
+
+def re_execute_with_namespace(argv, network=False):
+  """Re-execute as root so we can unshare resources."""
+  if os.geteuid() != 0:
+    cmd = [
+        'sudo',
+        'HOME=%s' % os.environ['HOME'],
+        'PATH=%s' % os.environ['PATH'],
+        '--',
+    ] + argv
+    os.execvp(cmd[0], cmd)
+  else:
+    cgroups.Cgroup.InitSystem()
+    namespaces.SimpleUnshare(net=not network, pid=True)
+    # We got our namespaces, so switch back to the user to run the tests.
+    gid = int(os.environ.pop('SUDO_GID'))
+    uid = int(os.environ.pop('SUDO_UID'))
+    user = os.environ.pop('SUDO_USER')
+    os.initgroups(user, gid)
+    os.setresgid(gid, gid, gid)
+    os.setresuid(uid, uid, uid)
+    os.environ['USER'] = user
+
+
+def re_execute_inside_chroot(argv):
+  """Re-execute the test wrapper inside the chroot."""
+  cmd = [
+      'cros_sdk',
+      '--',
+      os.path.join('..', '..', 'chromite', 'run_pytest'),
+  ]
+  if not cros_build_lib.IsInsideChroot():
+    os.execvp(cmd[0], cmd + argv)
+  else:
+    os.chdir(constants.CHROMITE_DIR)
+
+
+def ensure_chroot_exists():
+  """Ensure that a chroot exists for us to run tests in."""
+  chroot = os.path.join(constants.SOURCE_ROOT, constants.DEFAULT_CHROOT_DIR)
+  if not os.path.exists(chroot) and not cros_build_lib.IsInsideChroot():
+    cros_build_lib.run(['cros_sdk', '--create'])