# Copyright 2017 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Runner for Mobly test suites.

To create a test suite, call suite_runner.run_suite() with one or more
individual test classes. For example:

.. code-block:: python

    from mobly import suite_runner

    from my.test.lib import foo_test
    from my.test.lib import bar_test
    ...
    if __name__ == '__main__':
        suite_runner.run_suite(foo_test.FooTest, bar_test.BarTest)
"""

import argparse
import collections
import logging
import sys

from mobly import base_test
from mobly import config_parser
from mobly import signals
from mobly import test_runner


class Error(Exception):
    pass


def run_suite(test_classes, argv=None):
    """Executes multiple test classes as a suite.

    This is the default entry point for running a test suite script file
    directly.

    Args:
        test_classes: List of python classes containing Mobly tests.
        argv: A list that is then parsed as cli args. If None, defaults to cli
            input.
    """
    # Parse cli args.
    parser = argparse.ArgumentParser(description='Mobly Suite Executable.')
    parser.add_argument('-c',
                        '--config',
                        type=str,
                        required=True,
                        metavar='<PATH>',
                        help='Path to the test configuration file.')
    parser.add_argument(
        '--tests',
        '--test_case',
        nargs='+',
        type=str,
        metavar='[ClassA[.test_a] ClassB[.test_b] ...]',
        help='A list of test classes and optional tests to execute.')
    if not argv:
        argv = sys.argv[1:]
    args = parser.parse_args(argv)
    # Load test config file.
    test_configs = config_parser.load_test_config_file(args.config)

    # Check the classes that were passed in
    for test_class in test_classes:
        if not issubclass(test_class, base_test.BaseTestClass):
            logging.error(
                'Test class %s does not extend '
                'mobly.base_test.BaseTestClass', test_class)
            sys.exit(1)

    # Find the full list of tests to execute
    selected_tests = compute_selected_tests(test_classes, args.tests)

    # Execute the suite
    ok = True
    for config in test_configs:
        runner = test_runner.TestRunner(config.log_path, config.testbed_name)
        with runner.mobly_logger():
            for (test_class, tests) in selected_tests.items():
                runner.add_test_class(config, test_class, tests)
            try:
                runner.run()
                ok = runner.results.is_all_pass and ok
            except signals.TestAbortAll:
                pass
            except:
                logging.exception('Exception when executing %s.',
                                  config.testbed_name)
                ok = False
    if not ok:
        sys.exit(1)


def compute_selected_tests(test_classes, selected_tests):
    """Computes tests to run for each class from selector strings.

    This function transforms a list of selector strings (such as FooTest or
    FooTest.test_method_a) to a dict where keys are test_name classes, and
    values are lists of selected tests in those classes. None means all tests in
    that class are selected.

    Args:
        test_classes: list of strings, names of all the classes that are part
            of a suite.
        selected_tests: list of strings, list of tests to execute. If empty,
            all classes `test_classes` are selected. E.g.

            .. code-block:: python

                [
                    'FooTest',
                    'BarTest',
                    'BazTest.test_method_a',
                    'BazTest.test_method_b'
                ]

    Returns:
        dict: Identifiers for TestRunner. Keys are test class names; valures
            are lists of test names within class. E.g. the example in
            `selected_tests` would translate to:

            .. code-block:: python

                {
                    FooTest: None,
                    BarTest: None,
                    BazTest: ['test_method_a', 'test_method_b']
                }

            This dict is easy to consume for `TestRunner`.
    """
    class_to_tests = collections.OrderedDict()
    if not selected_tests:
        # No selection is needed; simply run all tests in all classes.
        for test_class in test_classes:
            class_to_tests[test_class] = None
        return class_to_tests

    # The user is selecting some tests to run. Parse the selectors.
    # Dict from test_name class name to list of tests to execute (or None for all
    # tests).
    test_class_name_to_tests = collections.OrderedDict()
    for test_name in selected_tests:
        if '.' in test_name:  # Has a test method
            (test_class_name, test_name) = test_name.split('.')
            if test_class_name not in test_class_name_to_tests:
                # Never seen this class before
                test_class_name_to_tests[test_class_name] = [test_name]
            elif test_class_name_to_tests[test_class_name] is None:
                # Already running all tests in this class, so ignore this extra
                # test.
                pass
            else:
                test_class_name_to_tests[test_class_name].append(test_name)
        else:  # No test method; run all tests in this class.
            test_class_name_to_tests[test_name] = None

    # Now transform class names to class objects.
    # Dict from test_name class name to instance.
    class_name_to_class = {cls.__name__: cls for cls in test_classes}
    for test_class_name, tests in test_class_name_to_tests.items():
        test_class = class_name_to_class.get(test_class_name)
        if not test_class:
            raise Error('Unknown test_name class %s' % test_class_name)
        class_to_tests[test_class] = tests

    return class_to_tests