#!/bin/bash

# Copyright (c) 2011 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.

CHROMITE_PATH=$(realpath "$(dirname "$0")/..")
IN_CHROOT="cros_sdk"
TIMEOUT="timeout -k 5m 20m"
CHROOT_CHROMITE=../../chromite

set -eu

# List all exceptions, with a token describing what's odd here.
# inside - inside the chroot
# skip - don't run this test (please comment on why)
declare -A special_tests
special_tests=(
  # Tests that need to run inside the chroot.
  ['cbuildbot/stages/test_stages_unittest.py']=inside
  ['cros/commands/cros_build_unittest.py']=inside
  ['cros/commands/lint_unittest.py']=inside
  ['lib/filetype_unittest.py']=inside
  ['lib/upgrade_table_unittest.py']=inside
  ['scripts/cros_list_modified_packages_unittest.py']=inside
  ['scripts/cros_mark_as_stable_unittest.py']=inside
  ['scripts/cros_mark_chrome_as_stable_unittest.py']=inside
  ['scripts/sync_package_status_unittest.py']=inside
  ['scripts/cros_portage_upgrade_unittest.py']=inside
  ['scripts/dep_tracker_unittest.py']=inside
  ['scripts/test_image_unittest.py']=inside
  ['scripts/upload_package_status_unittest.py']=inside

  # Tests that need to run outside the chroot.
  ['lib/cgroups_unittest.py']=outside

  # Tests that take >2 minutes to run.  All the slow tests are
  # disabled atm though ...
  #['scripts/cros_portage_upgrade_unittest.py']=skip
)

skip_quick_tests() {
  # Tests that require network can be really slow.
  special_tests['cbuildbot/manifest_version_unittest.py']=skip
  special_tests['cbuildbot/repository_unittest.py']=skip
  special_tests['cbuildbot/remote_try_unittest.py']=skip
  special_tests['lib/cros_build_lib_unittest.py']=skip
  special_tests['lib/gerrit_unittest.py']=skip
  special_tests['lib/patch_unittest.py']=skip

  # cgroups_unittest runs cros_sdk a lot, so is slow.
  special_tests['lib/cgroups_unittest.py']=skip
}

# Helper function to add failed logs/tests to be printed out later.
# $1 test that failed.
# $2 log file where we stored the output of the failed test.
append_failed_test() {
  echo "ERROR: Unittest $1 failed.  Log will be output at end of run!!!"

  cat - "$2" <<EOF >>"${LOGFILE}.err.${BASHPID}"

FAIL: Unittest $1 failed output:

EOF
}

# Wrapper to run unittest.  Hides output unless test fails.
# $1 test to run.  Must be in chromite/cbuildbot.
# $2 Is this a dry run?
# $3 Do we run network tests?
run_test() {
  local test=$1 dryrun=$2 network_tests=$3
  local log_file="${LOGFILE}.tmp.${BASHPID}"
  local special="${special_tests[${test}]:-}"
  local starttime="$(date +%s%N)"

  local network_flag=""
  if ${network_tests}; then
    network_flag="--network"
  fi

  if [[ "${special}" == "skip" ]]; then
    echo "Skipping unittest ${test}"
    return
  elif [[ "${special}" == "outside" && -f /etc/cros_chroot_version ]]; then
    echo "Skipping unittest ${test} because it must run outside the chroot"
    return
  elif [[ "${special}" == "inside" && ! -f /etc/cros_chroot_version ]]; then
    if ${SKIP_CHROOT_TESTS}; then
      echo "Skipping unittest ${test} because it must run inside the chroot"
      return
    else
      echo "Starting unittest ${test} inside the chroot"
      if ! ${dryrun}; then
        ${TIMEOUT} ${IN_CHROOT} -- python "${CHROOT_CHROMITE}/${test}" -v \
          ${network_flag} &> "${log_file}" \
          || append_failed_test "${test}" "${log_file}"
      fi
    fi
  else
    echo "Starting unittest ${test}"
    if ! ${dryrun}; then
      ${TIMEOUT} python "${CHROMITE_PATH}/${test}" -v ${network_flag} \
        &> "${log_file}" || append_failed_test "${test}" "${log_file}"
    fi
  fi

  if ! ${dryrun}; then
    local endtime="$(date +%s%N)"
    local duration=$(( (endtime - starttime) / 1000000 ))

    echo "Finished unittest ${test} (${duration} ms)"
  fi
  rm -f "${log_file}"
}

cleanup() {
  delayed_kill() {
    sleep 5
    kill -9 ${children[*]} &> /dev/null
  }

  echo "Cleaning up backgrounded jobs."
  # Graceful exit.
  kill -INT ${children[*]} &> /dev/null
  # Set of a hard kill timer after a while.
  delayed_kill &
  wait ${children[*]}
  show_logs
}

show_logs() {
  cat "${LOGFILE}".err.* > "${LOGFILE}" 2>/dev/null || :

  rm -f "${LOGFILE}".*
  if [[ -s ${LOGFILE} ]]; then
    cat "${LOGFILE}"
    echo
    echo
    echo "FAIL: The following tests failed:"
    sed -nre '/^FAIL:/s/^FAIL: Unittest (.*) failed output:/\1/p' "${LOGFILE}"
    rm -f "${LOGFILE}"
    exit 1
  fi

  rm -f "${LOGFILE}"
}

usage() {
  cat <<EOF
Usage: run_tests [options] [tests]

Run the specified tests.  If none are specified, we'll scan the
tree looking for tests to run and then only run the semi-fast ones.

You can add a .testignore file to a dir to disable scanning it.

Options:
  -q, --quick     Only run the really quick tests
  -n, --dry-run   Do everything but actually run the test
  --network       Run tests that depend on good network connectivity
  -l, --list      List all the available tests
  -h, --help      This screen
EOF

  if [[ $# -gt 0 ]]; then
    printf '\nerror: %s\n' "$*" >&2
    exit 1
  else
    exit 0
  fi
}

main() {
  # Parse args from the user first.
  local list=false
  local dryrun=false
  local user_tests=()
  local network_tests=false
  while [[ $# -gt 0 ]]; do
    case $1 in
    -h|--help)     usage;;
    -q|--quick)    skip_quick_tests;;
    -n|--dry-run)  dryrun=true;;
    -l|--list)     list=true;;
    --network)     network_tests=true;;
    -*)            usage "unknown option $1";;
    *)             user_tests+=( "$1" );;
    esac
    shift
  done
  if [[ ${#user_tests[@]} -gt 0 ]]; then
    set -- "${user_tests[@]}"
  fi

  # For some versions of 'sudo' (notably, the 'sudo' in the chroot at
  # the time of this writing), sudo -v will ask for a password whether
  # or not it's needed.  'sudo true' will do what we want.
  sudo true

  # Switch to CHROMITE_PATH, in case cwd is outside of the repo. This ensures
  # that "repo list" looks at the right repo, and sets up a consistent test
  # environment.
  cd "${CHROMITE_PATH}"

  # Make sure the index doesn't fall out of date.
  local k
  for k in "${!special_tests[@]}"; do
    if [[ ! -e ${k} ]]; then
      echo "ERROR: test '${k}' listed in \${special_tests} but does not exist."
      return 1
    fi
  done

  SKIP_CHROOT_TESTS=false
  if ! repo list 2>/dev/null | grep -q chromiumos-overlay; then
    echo "chromiumos-overlay is not present. Skipping chroot tests..."
    SKIP_CHROOT_TESTS=true

    # cgroups_unittest requires cros_sdk, so it doesn't work.
    special_tests['lib/cgroups_unittest.py']=skip
  fi

  # Default to running all tests.
  if [[ $# -eq 0 ]]; then
    # List all unit test scripts that match the given pattern.
    local prune_tests=(
      $(find "${CHROMITE_PATH}" -name .testignore \
          -printf '-path %h -prune -o ')
    )
    local all_tests=(
      $(find "${CHROMITE_PATH}" \
          ${prune_tests[*]} \
          -name '*_unittest.py' -printf '%P ')
    )
    set -- "${all_tests[@]}"
  fi

  if ${list}; then
    printf '%s\n' "$@" | sort
    exit 0
  fi

  # Now do the real code.
  LOGFILE="$(mktemp -t cbuildbot.run_tests.XXXXXX)"
  trap cleanup INT TERM

  local children=()
  for test in "$@"; do
    run_test ${test} ${dryrun} ${network_tests} &
    children+=( $! )
  done

  wait ${children[*]}
  trap - INT TERM

  show_logs
  if ! ${dryrun}; then
    echo "All tests succeeded!"
    if ! ${network_tests}; then
      echo "Note: Network tests skipped; use --network to run them."
    fi
  fi
}

main "$@"