# Copyright 2016 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Runs JUnit tests on a local machine."""

import dataclasses
import json
import logging
import multiprocessing
import os
import queue
import re
import subprocess
import sys
import tempfile
import threading
import time
import zipfile
from concurrent.futures import ThreadPoolExecutor

from devil.utils import cmd_helper
from py_utils import tempfile_ext
from pylib import constants
from pylib.base import base_test_result
from pylib.base import test_run
from pylib.constants import host_paths
from pylib.results import json_results

# Chosen after timing test runs of chrome_junit_tests with 7,16,32,
# and 64 workers in threadpool and different classes_per_job.
_MAX_TESTS_PER_JOB = 128

_FAILURE_TYPES = (
    base_test_result.ResultType.FAIL,
    base_test_result.ResultType.CRASH,
    base_test_result.ResultType.TIMEOUT,
)

# RegExp to detect logcat lines, e.g., 'I/AssetManager: not found'.
_LOGCAT_RE = re.compile(r' ?\d+\| (:?\d+\| )?[A-Z]/[\w\d_-]+:')

# Regex to detect start or failure of tests. Matches
# [ RUN      ] org.ui.ForeignSessionItemViewBinderUnitTest.test_phone[28]
# [ FAILED|CRASHED|TIMEOUT ] org.ui.ForeignBinderUnitTest.test_phone[28] (56 ms)
_TEST_START_RE = re.compile(r'.*\[\s+RUN\s+\]\s(.*)')
_TEST_FAILED_RE = re.compile(r'.*\[\s+(?:FAILED|CRASHED|TIMEOUT)\s+\]')
_TEST_FINISHED_RE = re.compile(r'.*\[\s+(?:OK|SKIPPED)\s+\]')


@dataclasses.dataclass
class _TestGroup:
    config: str
    methods_by_class: dict


@dataclasses.dataclass
class _Job:
    shard_id: int
    cmd: str
    timeout: int
    json_config: dict
    json_results_path: str


class LocalMachineJunitTestRun(test_run.TestRun):
    # override
    def TestPackage(self):
        return self._test_instance.suite

    # override
    def SetUp(self):
        pass

    def _GetFilterArgs(self):
        ret = []
        for test_filter in self._test_instance.test_filters:
            ret += ['--gtest-filter', test_filter]
        for test_filter_file in self._test_instance.test_filter_files:
            ret += ['--test-launcher-filter-file', test_filter_file]

        if self._test_instance.package_filter:
            ret += ['--package-filter', self._test_instance.package_filter]
        if self._test_instance.runner_filter:
            ret += ['--runner-filter', self._test_instance.runner_filter]

        return ret

    def _CreatePropertiesJar(self, temp_dir):
        # Create properties file for Robolectric test runners so they can find
        # the binary resources.
        properties_jar_path = os.path.join(temp_dir, 'properties.jar')
        resource_apk = self._test_instance.resource_apk
        with zipfile.ZipFile(properties_jar_path, 'w') as z:
            if resource_apk:
                z.writestr(
                    'com/android/tools/test_config.properties',
                    'android_resource_apk=%s\n' % resource_apk,
                )
            # These values must be kept in sync with
            # BaseRobolectricTestRunner.java.
            min_sdk = '29'
            max_sdk = '36'

            if self._test_instance.single_variant:
                sdk_string = max_sdk
            else:
                sdk_string = f"{min_sdk},{max_sdk}"
            props = [
                'application = android.app.Application',
                'sdk = %s' % sdk_string,
                (
                    'shadows = org.chromium.testing.local.'
                    'CustomShadowApplicationPackageManager,'
                    'org.chromium.testing.local.ShadowTextView'
                ),
            ]

            if not resource_apk:
                # Setting manifest = --none improves performance by avoiding
                # Robolectric having to scan for and parse a dummy manifest.
                props.append('manifest = --none')

            z.writestr('robolectric.properties', '\n'.join(props))
        return properties_jar_path

    def _CreateJvmArgsList(self, for_listing=False, allow_debugging=True):
        # Creates a list of jvm_args (robolectric, code coverage, etc...)
        jvm_args = [
            # JDK 17+ requires explicit opens for Robolectric reflection on
            # internal fields:
            # https://docs.oracle.com/en/java/javase/17/migrate/migrating-jdk-8-later-jdk-releases.html
            '--add-opens=java.base/java.io=ALL-UNNAMED',
            '--add-opens=java.base/java.lang=ALL-UNNAMED',
            '--add-opens=java.base/java.util=ALL-UNNAMED',
            '--add-opens=java.base/jdk.internal.access=ALL-UNNAMED',
            # Disable warning about mockito/bytebuddy dynamically adding an
            # agent.
            '-XX:+EnableDynamicAgentLoading',
            '-Drobolectric.dependency.dir=%s'
            % self._test_instance.robolectric_runtime_deps_dir,
            '-Ddir.source.root=%s' % constants.DIR_SOURCE_ROOT,
            # Use locally available sdk jars from 'robolectric.dependency.dir'
            '-Drobolectric.offline=true',
            '-Drobolectric.resourcesMode=binary',
            '-Drobolectric.logging=stdout',
            '-Drobolectric.enforceViewMethodsCalledOnMainThread=true',
            '-Djava.library.path=%s' % self._test_instance.native_libs_dir,
        ]
        if self._test_instance.run_disabled:
            jvm_args += ['-Dchromium.run_disabled=1']
        if self._test_instance.debug_socket and allow_debugging:
            jvm_args += [
                '-Dchromium.jdwp_active=true',
                (
                    '-agentlib:jdwp=transport=dt_socket'
                    ',server=y,suspend=y,address=%s'
                    % self._test_instance.debug_socket
                ),
            ]

        if self._test_instance.coverage_dir and not for_listing:
            if not os.path.exists(self._test_instance.coverage_dir):
                os.makedirs(self._test_instance.coverage_dir)
            elif not os.path.isdir(self._test_instance.coverage_dir):
                raise Exception(
                    '--coverage-dir takes a directory, not file path.'
                )
            # Jacoco supports concurrent processes using the same output file:
            # https://github.com/jacoco/jacoco/blob/6cd3f0bd8e348f8fba7bffec5225407151f1cc91/org.jacoco.agent.rt/src/org/jacoco/agent/rt/internal/output/FileOutput.java#L67
            # So no need to vary the output based on shard number.
            jacoco_coverage_file = os.path.join(
                self._test_instance.coverage_dir,
                '%s.exec' % self._test_instance.suite,
            )
            if self._test_instance.coverage_on_the_fly:
                jacoco_agent_path = os.path.join(
                    host_paths.DIR_SOURCE_ROOT,
                    'third_party',
                    'jacoco',
                    'cipd',
                    'lib',
                    'jacocoagent.jar',
                )

                # inclnolocationclasses is false to prevent no class def found
                # error.
                jacoco_args = (
                    '-javaagent:{}=destfile={},inclnolocationclasses=false'
                )
                jvm_args.append(
                    jacoco_args.format(jacoco_agent_path, jacoco_coverage_file)
                )
            else:
                jvm_args.append(
                    '-Djacoco-agent.destfile=%s' % jacoco_coverage_file
                )

        return jvm_args

    def _ChooseNumWorkers(self, num_jobs):
        if self._test_instance.debug_socket:
            num_workers = 1
        elif self._test_instance.shards is not None:
            num_workers = self._test_instance.shards
        else:
            num_workers = max(1, multiprocessing.cpu_count() // 2)
        return min(num_workers, num_jobs)

    def _ApplyExternalSharding(self, json_config):
        shard_index = self._test_instance.external_shard_index
        total_shards = self._test_instance.total_external_shards
        if total_shards <= 1:
            return json_config

        logging.info(
            'Using external sharding settings. This is shard %d/%d',
            shard_index,
            total_shards,
        )

        all_groups = GroupTests(json_config, _MAX_TESTS_PER_JOB)
        selected_groups = [
            g
            for i, g in enumerate(all_groups)
            if i % total_shards == shard_index
        ]

        new_configs = {}
        for group in selected_groups:
            if group.config not in new_configs:
                new_configs[group.config] = {}
            new_configs[group.config].update(group.methods_by_class)

        return {**json_config, 'configs': new_configs}

    @property
    def _wrapper_path(self):
        return os.path.join(
            constants.GetOutDirectory(),
            'bin',
            'helper',
            self._test_instance.suite,
        )

    def _QueryTestJsonConfig(
        self, temp_dir, allow_debugging=True, enable_shadow_allowlist=False
    ):
        json_config_path = os.path.join(temp_dir, 'main_test_config.json')
        cmd = [self._wrapper_path]
        # Allow debugging of test listing when run as:
        # "--wait-for-java-debugger --list-tests"
        jvm_args = self._CreateJvmArgsList(
            for_listing=True, allow_debugging=allow_debugging
        )
        if jvm_args:
            cmd += ['--jvm-args', '"%s"' % ' '.join(jvm_args)]
        cmd += ['--classpath', self._CreatePropertiesJar(temp_dir)]
        cmd += ['--list-tests', '--json-config', json_config_path]
        if enable_shadow_allowlist and self._test_instance.shadows_allowlist:
            cmd += [
                '--shadows-allowlist',
                self._test_instance.shadows_allowlist,
            ]
        cmd += self._GetFilterArgs()
        subprocess.run(cmd, check=True)
        with open(json_config_path, encoding='utf-8') as f:
            return json.load(f)

    def _MakeJob(
        self, shard_id, temp_dir, test_group, properties_jar_path, json_config
    ):
        json_results_path = os.path.join(temp_dir, f'results{shard_id}.json')
        job_json_config_path = os.path.join(temp_dir, f'config{shard_id}.json')
        job_json_config = json_config.copy()
        job_json_config['configs'] = {
            test_group.config: test_group.methods_by_class
        }
        with open(job_json_config_path, 'w', encoding='utf-8') as f:
            json.dump(job_json_config, f)

        cmd = [self._wrapper_path]
        cmd += ['--jvm-args', '"%s"' % ' '.join(self._CreateJvmArgsList())]
        cmd += ['--classpath', properties_jar_path]
        cmd += ['--json-results', json_results_path]
        cmd += ['--json-config', job_json_config_path]

        if self._test_instance.debug_socket:
            timeout = 999999
        else:
            # 20 seconds for process init,
            # 5 seconds per class,
            # 3 seconds per method.
            num_classes = len(test_group.methods_by_class)
            num_tests = sum(
                len(x) for x in test_group.methods_by_class.values()
            )
            timeout = 60 + 10 * num_classes + num_tests * 5
        return _Job(
            shard_id=shard_id,
            cmd=cmd,
            timeout=timeout,
            json_config=job_json_config,
            json_results_path=json_results_path,
        )

    def _GetJsonConfig(self):
        with tempfile_ext.NamedTemporaryDirectory() as temp_dir:
            return self._QueryTestJsonConfig(
                temp_dir, allow_debugging=False, enable_shadow_allowlist=True
            )

    # override
    def GetTestsForListing(self):
        json_config = self._GetJsonConfig()
        json_config = self._ApplyExternalSharding(json_config)
        ret = []
        for config in json_config['configs'].values():
            for class_name, methods in config.items():
                ret.extend(f'{class_name}.{method}' for method in methods)
        for config in json_config.get('disabled', {}).values():
            for class_name, methods in config.items():
                ret.extend(
                    f'{class_name}.{method} (disabled)' for method in methods
                )
        ret.sort()
        return ret

    # override
    def RunTests(self, results, raw_logs_fh=None):
        with tempfile_ext.NamedTemporaryDirectory() as temp_dir:
            self._RunTestsInternal(temp_dir, results, raw_logs_fh)

    def _RunTestsInternal(self, temp_dir, results, raw_logs_fh):
        if self._test_instance.json_config:
            with open(self._test_instance.json_config, encoding='utf-8') as f:
                json_config = json.load(f)
        else:
            try:
                json_config = self._GetJsonConfig()
            except (subprocess.CalledProcessError, IOError):
                results.append(
                    _MakeUnknownFailureResult('Filter matched no tests')
                )
                return

        # Disabled tests are nested in a top-level "disabled" key.
        # Merge them into the main "configs".
        if self._test_instance.run_disabled:
            for config, classes in json_config.get('disabled', {}).items():
                target_classes = json_config['configs'].setdefault(config, {})
                for class_name, methods in classes.items():
                    target_methods = target_classes.setdefault(class_name, [])
                    target_methods.extend(methods)

        json_config = self._ApplyExternalSharding(json_config)
        test_groups = GroupTests(json_config, _MAX_TESTS_PER_JOB)

        shard_list = list(range(len(test_groups)))
        shard_filter = self._test_instance.shard_filter
        if shard_filter:
            shard_list = [x for x in shard_list if x in shard_filter]

        if not shard_list:
            results.append(_MakeUnknownFailureResult('Invalid shard filter'))
            return

        num_workers = self._ChooseNumWorkers(len(shard_list))
        if shard_filter:
            logging.warning(
                'Running test shards: %s using %s concurrent process(es)',
                ', '.join(str(x) for x in shard_list),
                num_workers,
            )
        else:
            logging.warning(
                'Running tests with %d shard(s) using %s'
                ' concurrent process(es).',
                len(shard_list),
                num_workers,
            )

        properties_jar_path = self._CreatePropertiesJar(temp_dir)
        jobs = [
            self._MakeJob(
                i, temp_dir, test_groups[i], properties_jar_path, json_config
            )
            for i in shard_list
        ]

        show_logcat = logging.getLogger().isEnabledFor(logging.INFO)
        num_omitted_lines = 0
        failed_test_logs = {}
        log_lines = []
        current_test = None
        for line in RunCommandsAndSerializeOutput(
            jobs, num_workers, quiet=self._test_instance.quiet
        ):
            if raw_logs_fh:
                raw_logs_fh.write(line)

            if not self._test_instance.quiet:
                if show_logcat or not _LOGCAT_RE.match(line):
                    sys.stdout.write(line)
                else:
                    num_omitted_lines += 1

            # Collect log data between a test starting and the test failing.
            # There can be info after a test fails and before the next test
            # starts that we discard.
            test_start_match = _TEST_START_RE.match(line)
            if test_start_match:
                current_test = test_start_match.group(1)
                log_lines = [line]
            elif _TEST_FAILED_RE.match(line) and current_test:
                log_lines.append(line)
                failed_test_logs[current_test] = ''.join(log_lines)
                if self._test_instance.quiet:
                    for l in log_lines:
                        sys.stdout.write(l)
                current_test = None
            elif _TEST_FINISHED_RE.match(line) and current_test:
                current_test = None
            else:
                log_lines.append(line)

        if num_omitted_lines > 0:
            logging.critical('%d log lines omitted.', num_omitted_lines)
        sys.stdout.flush()

        results_list = []
        failed_jobs = []
        try:
            for job in jobs:
                with open(job.json_results_path, 'r', encoding='utf-8') as f:
                    parsed_results = json_results.ParseResultsFromJson(
                        json.loads(f.read())
                    )
                has_failed = False
                for r in parsed_results:
                    if r.GetType() in _FAILURE_TYPES:
                        has_failed = True
                        r.SetLog(
                            failed_test_logs.get(
                                r.GetName().replace('#', '.'), ''
                            )
                        )

                results_list += parsed_results
                if has_failed:
                    failed_jobs.append(job)
        except IOError:
            # In the case of a failure in the JUnit or Robolectric test runner
            # the output json file may never be written.
            results_list = [
                base_test_result.BaseTestResult(
                    'Test Runner Failure', base_test_result.ResultType.UNKNOWN
                )
            ]

        if failed_jobs and not self._test_instance.quiet:
            for job in failed_jobs:
                print(
                    f'To re-run failed shard {job.shard_id}, use --json-config '
                    'config.json, where config.json contains:'
                )
                print(json.dumps(job.json_config, indent=2))
                print()

            print(
                f'To re-run the {len(failed_jobs)} failed shard(s), use: '
                f'--shards {num_workers} --shard-filter',
                ','.join(str(j.shard_id) for j in failed_jobs),
            )

        test_run_results = base_test_result.TestRunResults()
        test_run_results.AddResults(results_list)

        num_actual = len(test_run_results.GetAll())
        num_expected = sum(
            len(methods)
            for classes in json_config['configs'].values()
            for methods in classes.values()
        )

        # Check that all tests actually ran that we expected to run.
        if num_actual != num_expected:
            sb = [
                f'Expected {num_expected} tests, but got only '
                f'{num_actual} results.',
                'Missing results for:',
            ]
            actual_test_names = {r.GetName() for r in test_run_results.GetAll()}
            for config, classes in json_config['configs'].items():
                for class_name, methods in classes.items():
                    for method in methods:
                        test_name = f'{class_name}#{method}'
                        if test_name not in actual_test_names:
                            sb.append(f'  {test_name} ({config})')
            if self._test_instance.run_disabled:
                sb.append(
                    'The missing tests could be due to not using '
                    'BaseJUnit4ClassRunner / BaseRobolectricTestRunner, since '
                    'these are required for --run-disabled.'
                )
            results.append(
                _MakeUnknownFailureResult('Not Enough Results', '\n'.join(sb))
            )

        if json_config.get('disabled') and not self._test_instance.run_disabled:
            num_disabled = sum(
                len(methods)
                for classes in json_config['disabled'].values()
                for methods in classes.values()
            )
            test_run.ShowDisabledTestsHint(count=num_disabled)

        results.append(test_run_results)

    # override
    def TearDown(self):
        pass


def GroupTests(json_config, max_per_job):
    """Groups tests that will be run on each shard.

    Args:
      json_config: The result from _QueryTestJsonConfig().
      max_per_job: Stop adding tests to a group once this limit has been passed.

    Return:
      Returns a list of _TestGroup.
    """
    ret = []
    for config, methods_by_class in json_config['configs'].items():
        size = 0
        group = {}
        for class_name, methods in methods_by_class.items():
            # There is some per-class overhead, so do not splits tests from one
            # class across multiple shards (unless configs differ).
            group[class_name] = methods
            size += len(methods)
            if size >= max_per_job:
                ret.append(_TestGroup(config, group))
                group = {}
                size = 0

        if group:
            ret.append(_TestGroup(config, group))

    # Put largest shards first to prevent long shards from being scheduled right
    # at the end.
    ret.sort(key=lambda x: -len(x.methods_by_class))
    return ret


def _MakeUnknownFailureResult(name, details=None):
    # TODO(504602174): These result types are not handled properly.
    results_list = [
        base_test_result.BaseTestResult(
            name, base_test_result.ResultType.UNKNOWN, log=details
        )
    ]
    test_run_results = base_test_result.TestRunResults()
    test_run_results.AddResults(results_list)
    logging.error('%s. Details: %s', name, details)
    return test_run_results


def _DumpJavaStacks(pid):
    jcmd = os.path.join(constants.JAVA_HOME, 'bin', 'jcmd')
    cmd = [jcmd, str(pid), 'Thread.print']
    result = subprocess.run(
        cmd, check=False, stdout=subprocess.PIPE, encoding='utf8'
    )
    if result.returncode:
        return 'Failed to dump stacks\n' + result.stdout
    return result.stdout


def RunCommandsAndSerializeOutput(jobs, num_workers, quiet=False):
    """Runs multiple commands in parallel and yields serialized output lines.

    Raises:
      TimeoutError: If timeout is exceeded.

    Yields:
      Command output.
    """
    assert jobs
    temp_files = [None]  # First shard is streamed directly to stdout.
    for _ in range(len(jobs) - 1):
        temp_files.append(tempfile.TemporaryFile(mode='w+t', encoding='utf-8'))

    if not quiet:
        yield '\n'
        yield f'Shard {jobs[0].shard_id} output:\n'

    timeout_dumps = {}

    def run_proc(idx):
        if idx == 0:
            s_out = subprocess.PIPE
            s_err = subprocess.STDOUT
        else:
            s_out = temp_files[idx]
            s_err = temp_files[idx]

        job = jobs[idx]
        proc = cmd_helper.Popen(
            job.cmd, stdout=s_out, stderr=s_err, env=getattr(job, 'env', None)
        )
        # Need to return process so that output can be displayed on stdout
        # in real time.
        if idx == 0:
            return proc

        try:
            proc.wait(timeout=job.timeout)
        except subprocess.TimeoutExpired:
            timeout_dumps[idx] = _DumpJavaStacks(proc.pid)
            proc.kill()

        # Not needed, but keeps pylint happy.
        return None

    with ThreadPoolExecutor(max_workers=num_workers) as pool:
        futures = [pool.submit(run_proc, idx=i) for i in range(len(jobs))]

        yield from _StreamFirstShardOutput(
            jobs[0], futures[0].result(), quiet=quiet
        )

        for i, job in enumerate(jobs[1:], 1):
            shard_id = job.shard_id
            # Shouldn't cause timeout as run_proc terminates the process with
            # a proc.wait().
            futures[i].result()
            f = temp_files[i]
            if not quiet:
                yield '\n'
                yield f'Shard {shard_id} output:\n'
            f.seek(0)
            for line in f.readlines():
                if quiet:
                    yield line
                else:
                    yield f'{shard_id:2}| {line}'
            f.close()

    # Output stacks
    if timeout_dumps:
        yield '\n'
        yield ('=' * 80) + '\n'
        yield '\nOne or more shards timed out.\n'
        yield ('=' * 80) + '\n'
        for i, dump in sorted(timeout_dumps.items()):
            job = jobs[i]
            yield (
                f'Shard {job.shard_id} timed out after {job.timeout} seconds.\n'
            )
            yield 'Thread dump:\n'
            yield dump
            yield '\n'

        raise cmd_helper.TimeoutError('Junit shards timed out.')


def _StreamFirstShardOutput(job, shard_proc, quiet=False):
    shard_id = job.shard_id
    # The following will be run from a thread to pump Shard 0 results, allowing
    # live output while allowing timeout.
    shard_queue = queue.Queue()

    def pump_stream_to_queue():
        for line in shard_proc.stdout:
            shard_queue.put(line)
        shard_queue.put(None)

    shard_0_pump = threading.Thread(target=pump_stream_to_queue)
    shard_0_pump.start()
    deadline = time.time() + job.timeout
    # Print the first process until timeout or completion.
    while shard_0_pump.is_alive():
        try:
            line = shard_queue.get(timeout=max(0, deadline - time.time()))
            if line is None:
                break
            if quiet:
                yield line
            else:
                yield f'{shard_id:2}| {line}'
        except queue.Empty:
            if time.time() > deadline:
                break

    # Output any remaining output from a timed-out first shard.
    shard_0_pump.join()
    while not shard_queue.empty():
        line = shard_queue.get()
        if line:
            if quiet:
                yield line
            else:
                yield f'{shard_id:2}| {line}'