# -*- coding: utf-8 -*-
# Copyright 2020 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Wrapper to execute pytest inside the chromite virtualenv."""

from __future__ import print_function

import argparse
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 cros_logging as logging
from chromite.lib import namespaces


def main(argv):
  parser = get_parser()
  opts, pytest_args = parser.parse_known_args()
  if opts.quick:
    if not cros_build_lib.IsInsideChroot() and opts.chroot:
      logging.warning('Tests start up faster when run from inside the chroot.')

  if opts.chroot:
    ensure_chroot_exists()
    re_execute_inside_chroot(argv)
  else:
    os.chdir(constants.CHROMITE_DIR)
    pytest_args += ['--no-chroot']

  # 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()

  if opts.quick:
    logging.info('Skipping test namespacing due to --quickstart.')
    # Default to running in a single process under --quickstart. User args can
    # still override this.
    pytest_args = ['-n', '0'] + pytest_args
  else:
    # Namespacing is enabled by default because tests may break each other or
    # interfere with parts of the running system if not isolated in a namespace.
    # Disabling namespaces is not recommended for general use.
    re_execute_with_namespace([sys.argv[0]] + argv)

  sys.exit(pytest.main(pytest_args))


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'])


def get_parser():
  """Build the parser for command line arguments."""
  parser = argparse.ArgumentParser(
      description=__doc__,
      epilog='To see the help output for pytest, run `pytest --help` inside '
      'the chroot.',
  )
  parser.add_argument(
      '--quickstart',
      dest='quick',
      action='store_true',
      help='Skip normal test sandboxing and namespacing for faster start up '
      'time.',
  )
  parser.add_argument(
      '--no-chroot',
      dest='chroot',
      action='store_false',
      help="Don't initialize or enter a chroot for the test invocation. May "
      'cause tests to unexpectedly fail!',
  )
  return parser