# Copyright 2019 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.

"""Test controller.

Handles all testing related functionality, it is not itself a test.
"""

import functools
import logging
import os
import string
import subprocess

from chromite.third_party.google.protobuf import json_format

from chromite.api import controller
from chromite.api import faux
from chromite.api import validate
from chromite.api.controller import controller_util
from chromite.api.gen.chromite.api import test_pb2
from chromite.api.gen.chromiumos import common_pb2
from chromite.api.gen.chromiumos.build.api import container_metadata_pb2
from chromite.api.metrics import deserialize_metrics_log
from chromite.lib import build_target_lib
from chromite.lib import chroot_lib
from chromite.lib import constants
from chromite.lib import cros_build_lib
from chromite.lib import goma_lib
from chromite.lib import metrics_lib
from chromite.lib import osutils
from chromite.lib import sysroot_lib
from chromite.lib.parser import package_info
from chromite.service import packages as packages_service
from chromite.service import test


@faux.empty_success
@faux.empty_completed_unsuccessfully_error
def DebugInfoTest(input_proto, _output_proto, config):
  """Run the debug info tests."""
  sysroot_path = input_proto.sysroot.path
  target_name = input_proto.sysroot.build_target.name

  if not sysroot_path:
    if target_name:
      sysroot_path = build_target_lib.get_default_sysroot_path(target_name)
    else:
      cros_build_lib.Die("The sysroot path or the sysroot's build target name "
                         'must be provided.')

  # We could get away with out this, but it's a cheap check.
  sysroot = sysroot_lib.Sysroot(sysroot_path)
  if not sysroot.Exists():
    cros_build_lib.Die('The provided sysroot does not exist.')

  if config.validate_only:
    return controller.RETURN_CODE_VALID_INPUT

  if test.DebugInfoTest(sysroot_path):
    return controller.RETURN_CODE_SUCCESS
  else:
    return controller.RETURN_CODE_COMPLETED_UNSUCCESSFULLY


def _BuildTargetUnitTestFailedResponse(_input_proto, output_proto, _config):
  """Add failed packages to a failed response."""
  packages = ['foo/bar', 'cat/pkg']
  for pkg in packages:
    pkg_info = package_info.parse(pkg)
    pkg_info_msg = output_proto.failed_packages.add()
    controller_util.serialize_package_info(pkg_info, pkg_info_msg)
    failed_pkg_data_msg = output_proto.failed_package_data.add()
    controller_util.serialize_package_info(pkg_info, failed_pkg_data_msg.name)
    failed_pkg_data_msg.log_path.path = '/path/to/%s/log' % pkg


@faux.empty_success
@faux.error(_BuildTargetUnitTestFailedResponse)
@validate.require_each('packages', ['category', 'package_name'])
@validate.validation_complete
@metrics_lib.collect_metrics
def BuildTargetUnitTest(input_proto, output_proto, _config):
  """Run a build target's ebuild unit tests."""
  # Method flags.
  # An empty sysroot means build packages was not run. This is used for
  # certain boards that need to use prebuilts (e.g. grunt's unittest-only).
  was_built = not input_proto.flags.empty_sysroot

  # Packages to be tested.
  packages_package_info = input_proto.packages
  packages = []
  for package_info_msg in packages_package_info:
    cpv = controller_util.PackageInfoToCPV(package_info_msg)
    packages.append(cpv.cp)

  # Skipped tests.
  blocklisted_package_info = input_proto.package_blocklist
  blocklist = []
  for package_info_msg in blocklisted_package_info:
    blocklist.append(controller_util.PackageInfoToString(package_info_msg))

  # Allow call to filter out non-cros_workon packages from the input packages.
  filter_only_cros_workon = input_proto.flags.filter_only_cros_workon

  # Allow call to succeed if no tests were found.
  testable_packages_optional = input_proto.flags.testable_packages_optional

  build_target = controller_util.ParseBuildTarget(input_proto.build_target)

  code_coverage = input_proto.flags.code_coverage

  sysroot = sysroot_lib.Sysroot(build_target.root)

  result = test.BuildTargetUnitTest(
      build_target,
      packages=packages,
      blocklist=blocklist,
      was_built=was_built,
      code_coverage=code_coverage,
      testable_packages_optional=testable_packages_optional,
      filter_only_cros_workon=filter_only_cros_workon)

  if not result.success:
    # Record all failed packages and retrieve log locations.
    controller_util.retrieve_package_log_paths(result.failed_pkgs,
                                               output_proto, sysroot)
    if result.failed_pkgs:
      return controller.RETURN_CODE_UNSUCCESSFUL_RESPONSE_AVAILABLE
    else:
      return controller.RETURN_CODE_COMPLETED_UNSUCCESSFULLY

  deserialize_metrics_log(output_proto.events, prefix=build_target.name)


SRC_DIR = os.path.join(constants.SOURCE_ROOT, 'src')
PLATFORM_DEV_DIR = os.path.join(SRC_DIR, 'platform/dev')
TEST_SERVICE_DIR = os.path.join(PLATFORM_DEV_DIR, 'src/chromiumos/test')


def _BuildTestServiceContainersResponse(input_proto, output_proto, _config):
  """Fake success response"""
  # pylint: disable=unused-argument
  output_proto.results.append(test_pb2.TestServiceContainerBuildResult(
      success=test_pb2.TestServiceContainerBuildResult.Success()
  ))


def _BuildTestServiceContainersFailedResponse(
    _input_proto, output_proto, _config):
  """Fake failure response"""

  # pylint: disable=unused-argument
  output_proto.results.append(test_pb2.TestServiceContainerBuildResult(
      failure=test_pb2.TestServiceContainerBuildResult.Failure(
          error_message='fake error'
      )
  ))


@validate.constraint('valid docker tag')
def _ValidDockerTag(tag):
  """Check that a string meets requirements for Docker tag naming."""
  # Tags can't start with period or dash
  if tag[0] in '.-':
    return "tag can't begin with '.' or '-'"

  # Tags can only consist of [a-zA-Z0-9-_.]
  allowed_chars = set(string.ascii_letters+string.digits+'-_.')
  invalid_chars = set(tag) - allowed_chars
  if invalid_chars:
    return f'saw one or more invalid characters: [{"".join(invalid_chars)}]'

  # Finally, max tag length is 128 characters
  if len(tag) > 128:
    return 'maximum tag length is 128 characters'


@validate.constraint('valid docker label key')
def _ValidDockerLabelKey(key):
  """Check that a string meets requirements for Docker tag naming."""

  # Label keys should start and end with a lowercase letter
  lowercase = set(string.ascii_lowercase)
  if not (key[0] in lowercase and key[-1] in lowercase):
    return "label key doesn't start and end with lowercase letter"

  # Label keys can have lower-case alphanumeric characters, period and dash
  allowed_chars = set(string.ascii_lowercase+string.digits+'-.')
  invalid_chars = set(key) - allowed_chars
  if invalid_chars:
    return f'saw one or more invalid characters: [{"".join(invalid_chars)}]'

  # Repeated . and - aren't allowed
  for char in '.-':
    if char * 2 in key:
      return f"'{char}' can\'t be repeated in label key"


@faux.success(_BuildTestServiceContainersResponse)
@faux.error(_BuildTestServiceContainersFailedResponse)
@validate.require('build_target.name')
@validate.require('chroot.path')
@validate.check_constraint('tags', _ValidDockerTag)
@validate.check_constraint('labels', _ValidDockerLabelKey)
@validate.validation_complete
def BuildTestServiceContainers(
    input_proto: test_pb2.BuildTestServiceContainersRequest,
    output_proto: test_pb2.BuildTestServiceContainersResponse, _config):
  """Builds docker containers for all test services and pushes them to gcr.io"""
  build_target = controller_util.ParseBuildTarget(input_proto.build_target)
  chroot = controller_util.ParseChroot(input_proto.chroot)
  sysroot = sysroot_lib.Sysroot(build_target.root)

  tags = ','.join(input_proto.tags)
  labels = (
      f'{key}={value}' for key, value in input_proto.labels.items()
  )

  build_script = os.path.join(
      TEST_SERVICE_DIR, 'python/src/docker_libs/cli/build-dockerimages.py')
  human_name = 'Service Builder'

  with osutils.TempDir(prefix='test_container') as tempdir:
    result_file = 'metadata.jsonpb'
    output_path = os.path.join(tempdir, result_file)
    # Note that we use an output file instead of stdout to avoid any issues
    # with maintaining stdout hygiene.  Stdout and stderr are combined to
    # form the error log in response to any errors.
    cmd = [build_script, chroot.path, sysroot.path]

    if input_proto.HasField('repository'):
      cmd += ['--host', input_proto.repository.hostname]
      cmd += ['--project', input_proto.repository.project]

    cmd += ['--tags', tags]
    cmd += ['--output', output_path]

    # Translate generator to comma separated string.
    ct_labels = ','.join(labels)
    cmd += ['--labels', ct_labels]
    cmd += ['--build_all']
    cmd += ['--upload']

    cmd_result = cros_build_lib.run(cmd, check=False,
                                    stderr=subprocess.STDOUT,
                                    stdout=True)

    if cmd_result.returncode != 0:
      # When failing, just record a fail response with the builder name.
      logging.debug('%s build failed.\nStdout:\n%s\nStderr:\n%s',
                    human_name, cmd_result.stdout, cmd_result.stderr)
      result = test_pb2.TestServiceContainerBuildResult()
      result.name = human_name
      image_info = container_metadata_pb2.ContainerImageInfo()
      result.failure.CopyFrom(
          test_pb2.TestServiceContainerBuildResult.Failure(
              error_message=cmd_result.stdout
          )
      )
      output_proto.results.append(result)

    else:
      logging.debug('%s build succeeded.\nStdout:\n%s\nStderr:\n%s',
                    human_name, cmd_result.stdout, cmd_result.stderr)
      files = os.listdir(tempdir)
      # Iterate through the tempdir to output metadata files.
      for file in files:
        if result_file in file:
          output_path = os.path.join(tempdir, file)

          # build-dockerimages.py will append the service name to outputfile
          # with an underscore.
          human_name = file.split('_')[-1]

          result = test_pb2.TestServiceContainerBuildResult()
          result.name = human_name
          image_info = container_metadata_pb2.ContainerImageInfo()
          json_format.Parse(osutils.ReadFile(output_path), image_info)
          result.success.CopyFrom(
              test_pb2.TestServiceContainerBuildResult.Success(
                  image_info=image_info
              )
          )
          output_proto.results.append(result)


@faux.empty_success
@faux.empty_completed_unsuccessfully_error
@validate.validation_complete
def ChromiteUnitTest(_input_proto, _output_proto, _config):
  """Run the chromite unit tests."""
  if test.ChromiteUnitTest():
    return controller.RETURN_CODE_SUCCESS
  else:
    return controller.RETURN_CODE_COMPLETED_UNSUCCESSFULLY


@faux.empty_success
@faux.empty_completed_unsuccessfully_error
@validate.validation_complete
def ChromitePytest(_input_proto, _output_proto, _config):
  """Run the chromite unit tests."""
  # TODO(vapier): Delete this stub.
  return controller.RETURN_CODE_SUCCESS


@faux.empty_success
@faux.empty_completed_unsuccessfully_error
@validate.validation_complete
def RulesCrosUnitTest(_input_proto, _output_proto, _config):
  """Run the rules_cros unit tests."""
  if test.RulesCrosUnitTest():
    return controller.RETURN_CODE_SUCCESS
  else:
    return controller.RETURN_CODE_COMPLETED_UNSUCCESSFULLY


@faux.all_empty
@validate.require('sysroot.path', 'sysroot.build_target.name', 'chrome_root')
@validate.validation_complete
def SimpleChromeWorkflowTest(input_proto, _output_proto, _config):
  """Run SimpleChromeWorkflow tests."""
  if input_proto.goma_config.goma_dir:
    chromeos_goma_dir = input_proto.goma_config.chromeos_goma_dir or None
    goma = goma_lib.Goma(
        input_proto.goma_config.goma_dir,
        input_proto.goma_config.goma_client_json,
        stage_name='BuildApiTestSimpleChrome',
        chromeos_goma_dir=chromeos_goma_dir)
  else:
    goma = None
  return test.SimpleChromeWorkflowTest(input_proto.sysroot.path,
                                       input_proto.sysroot.build_target.name,
                                       input_proto.chrome_root,
                                       goma)


@faux.all_empty
@validate.require('build_target.name', 'vm_path.path', 'test_harness',
                  'vm_tests')
@validate.validation_complete
def VmTest(input_proto, _output_proto, _config):
  """Run VM tests."""
  build_target_name = input_proto.build_target.name
  vm_path = input_proto.vm_path.path

  test_harness = input_proto.test_harness

  vm_tests = input_proto.vm_tests

  cmd = ['cros_run_test', '--debug', '--no-display', '--copy-on-write',
         '--board', build_target_name, '--image-path', vm_path,
         '--%s' % test_pb2.VmTestRequest.TestHarness.Name(test_harness).lower()]
  cmd.extend(vm_test.pattern for vm_test in vm_tests)

  if input_proto.ssh_options.port:
    cmd.extend(['--ssh-port', str(input_proto.ssh_options.port)])

  if input_proto.ssh_options.private_key_path:
    cmd.extend(['--private-key', input_proto.ssh_options.private_key_path.path])

  # TODO(evanhernandez): Find a nice way to pass test_that-args through
  # the build API. Or obviate them.
  if test_harness == test_pb2.VmTestRequest.AUTOTEST:
    cmd.append('--test_that-args=--allow-chrome-crashes')

  with osutils.TempDir(prefix='vm-test-results.') as results_dir:
    cmd.extend(['--results-dir', results_dir])
    cros_build_lib.run(cmd, kill_timeout=10 * 60)


@faux.all_empty
@validate.validation_complete
def CrosSigningTest(_input_proto, _output_proto, _config):
  """Run the cros-signing unit tests."""
  test_runner = os.path.join(constants.SOURCE_ROOT, 'cros-signing', 'signer',
                             'run_tests.py')
  result = cros_build_lib.run([test_runner], check=False)

  return result.returncode


def GetArtifacts(in_proto: common_pb2.ArtifactsByService.Test,
                 chroot: chroot_lib.Chroot, sysroot_class: sysroot_lib.Sysroot,
                 build_target: build_target_lib.BuildTarget,
                 output_dir: str) -> list:
  """Builds and copies test artifacts to specified output_dir.

  Copies test artifacts to output_dir, returning a list of (output_dir: str)
  paths to the desired files.

  Args:
    in_proto: Proto request defining reqs.
    chroot: The chroot class used for these artifacts.
    sysroot_class: The sysroot class used for these artifacts.
    build_target: The build target used for these artifacts.
    output_dir: The path to write artifacts to.

  Returns:
    A list of dictionary mappings of ArtifactType to list of paths.
  """
  generated = []

  artifact_types = {
      in_proto.ArtifactType.CODE_COVERAGE_LLVM_JSON:
          test.BundleCodeCoverageLlvmJson,
      in_proto.ArtifactType.HWQUAL: functools.partial(
          test.BundleHwqualTarball,
          build_target.name, packages_service.determine_full_version()),
  }

  for output_artifact in in_proto.output_artifacts:
    for artifact_type, func in artifact_types.items():
      if artifact_type in output_artifact.artifact_types:
        paths = func(chroot, sysroot_class, output_dir)
        if paths:
          generated.append({
              'paths': [paths] if isinstance(paths, str) else paths,
              'type': artifact_type,
          })

  return generated