# Copyright 2014 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Script to mount a built image and run tests on it."""

import contextlib
import os
import unittest

from chromite.lib import commandline
from chromite.lib import constants
from chromite.lib import image_lib
from chromite.lib import image_test_lib
from chromite.lib import osutils
from chromite.lib import path_util


def ParseArgs(args):
    """Return parsed commandline arguments."""

    parser = commandline.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--test_results_root",
        type="str_path",
        help="Directory to store test results",
    )
    parser.add_argument("--board", type=str, help="Board (wolf, beaglebone...)")
    parser.add_argument(
        "image",
        type="str_path",
        help="Image directory or file to test",
    )

    parser.add_argument(
        "-l",
        "--list",
        default=False,
        action="store_true",
        help="List all the available tests",
    )
    parser.add_bool_argument(
        "--bind-mount",
        default=False,
        enabled_desc="bind mount the image path as ROOT-A",
        disabled_desc="mount ROOT-A and STATE from the image file",
    )
    parser.add_argument(
        "tests",
        nargs="*",
        metavar="test",
        help="Specific tests to run (default runs all)",
    )

    opts = parser.parse_args(args)
    opts.Freeze()
    return opts


def FindImage(image_path):
    """Return the path to the image file.

    Args:
        image_path: A path to the image file, or a directory containing the base
            image.

    Returns:
        ImageFileAndMountScripts containing absolute paths to the image,
        the mount and umount invocation commands
    """

    if os.path.isdir(image_path):
        # Assume base image.
        image_file = os.path.join(
            image_path, constants.BASE_IMAGE_NAME + ".bin"
        )
        if not os.path.exists(image_file):
            raise ValueError("Cannot find base image %s" % image_file)
    elif os.path.isfile(image_path):
        image_file = image_path
    else:
        raise ValueError("%s is neither a directory nor a file" % image_path)

    return image_file


def main(args):
    opts = ParseArgs(args)

    # Build up test suites.
    loader = unittest.TestLoader()
    loader.suiteClass = image_test_lib.ImageTestSuite
    # We use a different prefix here so that unittest DO NOT pick up the
    # image tests automatically because they depend on a proper environment.
    loader.testMethodPrefix = "Test"
    tests_namespace = "chromite.cros.test.image_test"
    if opts.tests:
        tests = ["%s.%s" % (tests_namespace, x) for x in opts.tests]
    else:
        tests = (tests_namespace,)
    all_tests = loader.loadTestsFromNames(tests)

    # If they just want to see the lists of tests, show them now.
    if opts.list:

        def _WalkSuite(suite):
            for test in suite:
                if isinstance(test, unittest.BaseTestSuite):
                    for result in _WalkSuite(test):
                        yield result
                else:
                    yield (
                        test.id()[len(tests_namespace) + 1 :],
                        test.shortDescription() or "",
                    )

        test_list = list(_WalkSuite(all_tests))
        maxlen = max(len(x[0]) for x in test_list)
        for name, desc in test_list:
            print("%-*s  %s" % (maxlen, name, desc))
        return

    # Run them in the image directory.
    runner = image_test_lib.ImageTestRunner()
    runner.SetBoard(opts.board)
    runner.SetResultDir(opts.test_results_root)
    tmp_in_chroot = path_util.FromChrootPath("/tmp")

    with contextlib.ExitStack() as stack:
        temp_dir = stack.enter_context(osutils.TempDir(base_dir=tmp_in_chroot))

        if opts.bind_mount:
            runner.SetImageType(image_test_lib.ImageType.BIND_MOUNT)
            mount_path = temp_dir + "/dir-" + constants.PART_ROOT_A
            stack.enter_context(
                osutils.MountDirContext(
                    opts.image, mount_path, mount_opts=("bind", "ro")
                )
            )
        else:
            runner.SetImageType(image_test_lib.ImageType.BIN_FILE)
            image_file = FindImage(opts.image)
            stack.enter_context(
                image_lib.LoopbackPartitions(
                    image_file,
                    temp_dir,
                    (constants.PART_ROOT_A, constants.PART_STATE),
                )
            )

        with osutils.ChdirContext(temp_dir):
            result = runner.run(all_tests)

    if result and not result.wasSuccessful():
        return 1
    return 0