#!/usr/bin/env python3
############################################################################
# Copyright (C) 2006 The Regents of the University of California.
# Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
# Written by Christopher J. Morrone <morrone2@llnl.gov>
# CODE-OCEC-09-009. All rights reserved.
#
# This file is part of Slurm, a resource management program.
# For details, see <https://slurm.schedmd.com/>.
# Please also read the supplied file: DISCLAIMER.
#
# Slurm is free software; you can redistribute it and/or modify it under
# the terms of the GNU General Public License as published by the Free
# Software Foundation; either version 2 of the License, or (at your option)
# any later version.
#
# Slurm is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
# FOR A PARTICULAR PURPOSE.  See the GNU General Public License for more
# details.
#
# You should have received a copy of the GNU General Public License along
# with Slurm; if not, write to the Free Software Foundation, Inc.,
# 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA.
############################################################################

"""This script makes it easier to run the Slurm expect test scripts."""

from __future__ import print_function
import json
import os
import re
import sys
import time
import signal
from optparse import OptionParser
from optparse import OptionValueError
from subprocess import Popen


def main(argv=None):
    # "tests" is a list containing tuples of length 3 of the form
    # (test major number, test minor number, test filename)
    tests = []
    failed_tests = []
    passed_tests = []
    skipped_tests = []
    begin = (1, 1)
    abort = False

    # Handle command line parameters
    if argv is None:
        argv = sys.argv

    parser = OptionParser()
    parser.add_option(
        "-t",
        "--time-individual",
        action="store_true",
        dest="time_individual",
        default=False,
    )
    parser.add_option(
        "-e",
        "--exclude",
        type="string",
        dest="exclude_tests",
        action="callback",
        callback=test_parser,
        help="comma or space separated string of tests to skip",
    )
    parser.add_option(
        "-i",
        "--include",
        type="string",
        dest="include_tests",
        action="callback",
        callback=test_parser,
        help="comma or space separated string of tests to include",
    )
    parser.add_option("-k", "--keep-logs", action="store_true", default=False)
    parser.add_option("-s", "--stop-on-first-fail", action="store_true", default=False)
    parser.add_option(
        "-b",
        "--begin-from-test",
        type="string",
        dest="begin_from_test",
        action="callback",
        callback=test_parser,
    )
    parser.add_option(
        "-f",
        "--results-file",
        type="string",
        help="write json result to specified file name",
    )

    (options, args) = parser.parse_args(args=argv)

    # Sanity check
    if not os.path.isfile("globals"):
        print('ERROR: "globals" not here as needed', file=sys.stderr)
        return -1

        # Clear any environment variables that could break the tests.
        # Cray sets some squeue format options that break tests
        del os.environ["SQUEUE_ALL"]
        del os.environ["SQUEUE_SORT"]
        del os.environ["SQUEUE_FORMAT"]
        del os.environ["SQUEUE_FORMAT2"]

    # Read the current working directory and build a sorted list
    # of the available tests.
    test_re = re.compile(r"test(\d+)\.(\d+)$")
    for filename in os.listdir("."):
        match = test_re.match(filename)
        if match:
            major = int(match.group(1))
            minor = int(match.group(2))
            if not test_in_list(major, minor, options.exclude_tests) and (
                not options.include_tests
                or test_in_list(major, minor, options.include_tests)
            ):
                tests.append((major, minor, filename))
    if not tests:
        print(
            "ERROR: no test files found in current working directory", file=sys.stderr
        )
        return -1
    # sort by major, minor
    tests.sort(key=lambda t: (t[0], t[1]))

    # Set begin value
    if options.begin_from_test is not None:
        begin = options.begin_from_test[0]

    # Now run the tests
    start_time = time.time()
    test_env = os.environ.copy()
    if options.stop_on_first_fail:
        test_env["SLURM_TESTSUITE_CLEANUP_ON_FAILURE"] = "false"
    else:
        test_env["SLURM_TESTSUITE_CLEANUP_ON_FAILURE"] = "true"
    print("Started:", time.asctime(time.localtime(start_time)), file=sys.stdout)
    sys.stdout.flush()
    results_list = []
    for test in tests:
        if begin[0] > test[0] or (begin[0] == test[0] and begin[1] > test[1]):
            continue
        test_id = f"{test[0]}.{test[1]}"
        sys.stdout.write(f"Running test {test_id} ")
        sys.stdout.flush()
        test_dict = {}
        test_dict["id"] = test_id
        testlog_name = f"test{test_id}.log"
        try:
            os.remove(testlog_name + ".failed")
        except Exception:
            pass

        if os.path.exists(testlog_name):
            os.remove(testlog_name)
        testlog = open(testlog_name, "w+")

        if options.time_individual:
            t1 = time.time()
            test_dict["start_time"] = float("%.03f" % t1)

        try:
            child = Popen(
                ("expect", test[2]),
                shell=False,
                env=test_env,
                stdout=testlog,
                stderr=testlog,
            )
            retcode = child.wait()
        except KeyboardInterrupt:
            child.send_signal(signal.SIGINT)
            retcode = child.wait()
            abort = True

        if options.time_individual:
            t2 = time.time()
            minutes = int(int(t2 - t1) / 60)
            seconds = (int(t2 - t1)) % 60
            if minutes > 0:
                sys.stdout.write("%d min " % (minutes))
            sys.stdout.write("%.2f sec " % (seconds))
            test_dict["duration"] = float("%.03f" % (t2 - t1))

        if retcode == 0:
            status = "pass"
        elif retcode > 127:
            status = "skip"
        else:
            status = "fail"

        test_dict["status"] = status

        # Determine the reason if requesting a json results file
        if status != "pass" and options.results_file:
            testlog.flush()
            testlog.seek(0)
            test_output = testlog.read()

            sections = [s for s in test_output.split("=" * 78 + "\n")]
            # header = sections[1]
            body = sections[2]
            # footer = "".join(sections[3:])

            fatals = re.findall(
                r"(?ms)\[[^\]]+\][ \[]+Fatal[ \]:]+(.*?) \(fail[^\)]+\)$", body
            )
            errors = re.findall(
                r"(?ms)\[[^\]]+\][ \[]+Error[ \]:]+(.*?) \(subfail[^\)]+\)$", body
            )
            warnings = re.findall(
                r"(?ms)\[[^\]]+\][ \[]+Warning[ \]:]+((?:(?!Warning).)*) \((?:sub)?skip[^\)]+\)$",
                body,
            )
            if fatals:
                test_dict["reason"] = fatals[0]
            elif errors:
                test_dict["reason"] = errors[0]
            elif warnings:
                test_dict["reason"] = warnings[0]

        results_list.append(test_dict)

        testlog.close()

        if status == "pass":
            passed_tests.append(test)
            sys.stdout.write("\n")
            if not options.keep_logs:
                try:
                    os.remove(testlog_name)
                except IOError as e:
                    print(
                        "ERROR failed to close %s %s" % (testlog_name, e),
                        file=sys.stederr,
                    )
        elif status == "skip":
            skipped_tests.append(test)
            sys.stdout.write("SKIPPED\n")
            if not options.keep_logs:
                try:
                    os.remove(testlog_name)
                except IOError as e:
                    print(
                        "ERROR failed to close %s %s" % (testlog_name, e),
                        file=sys.stederr,
                    )
        else:
            failed_tests.append(test)
            os.rename(testlog_name, testlog_name + ".failed")
            sys.stdout.write("FAILED!\n")
            if options.stop_on_first_fail:
                break
        sys.stdout.flush()

        if abort:
            sys.stdout.write("\nRegression interrupted!\n")
            break

    end_time = time.time()
    print("Ended:", time.asctime(time.localtime(end_time)), file=sys.stdout)
    print(
        "\nTestsuite ran for %d minutes %d seconds"
        % ((end_time - start_time) / 60, (end_time - start_time) % 60),
        file=sys.stdout,
    )

    if options.results_file:
        with open(options.results_file, "w") as results_file:
            json.dump(results_list, results_file)

    print("Completions  :", len(passed_tests), file=sys.stdout)
    print("Failures     :", len(failed_tests), file=sys.stdout)
    print("Skipped      :", len(skipped_tests), file=sys.stdout)
    if len(failed_tests) > 0:
        print("Failed tests : ", file=sys.stdout)
        first = True
        for test in failed_tests:
            if first:
                first = False
            else:
                sys.stdout.write(",")
            sys.stdout.write("%d.%d" % (test[0], test[1]))
        sys.stdout.write("\n")
        sys.stdout.flush()

    if abort:
        print("INCOMPLETE", file=sys.stdout)

    if len(failed_tests) > 0:
        return 1


def test_in_list(major, minor, test_list):
    """Test for whether a test numbered major.minor is in test_list.

    "major" and "minor" must be integers.  "test_list" is a list of
    tuples, each tuple representing one test.  The tuples are of the
    form:

       (major, minor, filename)

    Returns True if the test is in the list, and False otherwise.
    """

    if not test_list:
        return False
    for test in test_list:
        if (test[0] == "*" or test[0] == major) and (
            test[1] == "*" or test[1] == minor
        ):
            return True
    return False


def test_parser(option, opt_str, value, parser):
    """Option callback function for the optparse.OptionParser class.

    Will take a string representing one or more test names and append
    a tuple representing the test into a list in the options's destination
    variable.

    A string representing test names must patch the regular expression
    named "test_re" below.  Some examples of exceptable options are:

        '1.5'
        'test9.8'
        '2.6 test3.1 14.2'
        '3.4,6.7,8.3'
        '1.*'
        '*.2'
        '1.*,3.8,9.2'

    Raises OptionValueError on error.
    """

    # Initialize the option's destination array, if is does not already exist.
    if not hasattr(parser.values, option.dest):
        setattr(parser.values, option.dest, [])
    if getattr(parser.values, option.dest) is None:
        setattr(parser.values, option.dest, [])

    # Get a pointer to the option's destination array.
    dest = getattr(parser.values, option.dest)

    # Split the user's option string into a series of tuples that represent
    # each test, and add each tuple to the destination array.
    splitter = re.compile(r"[,\s]+")
    val = splitter.split(value)
    test_re = re.compile(r"(test)?((\d+)|\*)[\._]((\d+)|\*)$")
    for v in val:
        m = test_re.match(v)
        if not m:
            raise OptionValueError
        major = m.group(2)
        if major != "*":
            major = int(major)
        minor = m.group(4)
        if minor != "*":
            minor = int(minor)
        dest.append((major, minor))


if __name__ == "__main__":
    sys.exit(main())