#!/usr/bin/env python
# Copyright 2022 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.
"""Detect flakiness in the Skia Gold based pixel tests.

This script runs the specified Skia Gold pixel tests multiple times and compares
screenshots generated by test runs. The flakiness is detected if pixel test
code generates different screenshots in different iterations.

Because screenshots are compared through MD5, this script should only check
the pixel tests that use precise matching.

This script only checks whether the screenshots under the same names change in
different iterations. This script does NOT check whether screenshots are
expected. Therefore, please ensure the screenshots are correct before running
this script.

During execution, this script creates directories for temporary data. Those
directories' names contain special characters to ensure uniqueness. This
script guarantees to delete those directories at the end of execution.
Users can control the paths of those temporary directories via the option
--root_dir.

* Example usages:

./tools/pixel_test/check_pixel_test_flakiness.py --gtest_filter=\
DemoAshPixelDiffTest.VerifyTopLevelWidgets --test_target=out/debug\
/ash_pixeltests --output_dir=../var

The command above should be executed at the chromium source directory whose
absolute file path looks like .../chromium/src. This command checks
DemoAshPixelDiffTest.VerifyTopLevelWidgets by running ash_pixeltests under the
directory .../chromium/src/out/debug. If flakiness is detected, the flaky test's
screenshots are saved under .../chromium/var. If the directory specified by
--output_dir does not exist and meanwhile the flakiness is detected, the script
will create one. If the --output_dir is not specified, the flaky test's
screenshots are not saved.

./tools/pixel_test/check_pixel_test_flakiness.py --gtest_filter=\
DemoAshPixelDiffTest.VerifyTopLevelWidgets --test_target=out/debug/\
ash_pixeltests --root_dir=../.. --output_dir=var

The command above is similar to the previous one. But difference is that this
command uses the option --root_dir to designate the root path for outputs
(including the temporary data and the saved screenshots when flakiness is
detected). In this example, the absolute path of the output directory is
.../chromium/../var rather than .../chromium/var.

./tools/pixel_test/check_pixel_test_flakiness.py --gtest_filter=\
*PersonalizationAppIntegrationPixel* --test_target=out/debug/browser_tests
--root_dir=/tmp/skia_gold --output_dir=var --browser-ui-tests-verify-pixels
--enable-pixel-output-in-tests

Finally, the above command runs the browser_tests target and adds extra
arguments necessary for experimental browser pixel tests to run properly.

* options:

--test_target: it specifies the path to the executable file of pixel tests. It
is a relative file path from the current working directory. The test target can
be any test executable based on Skia Gold.

--root_dir: it specifies the root path for outputs (including the temporary data
 and the saved screenshots when flakiness is detected). It is a relative file
 path from the current working directory.

--log_mode: its value can only be 'none', 'error_only' and 'all'. 'none' means
that the log generated by gTest runs does not show; 'error_only' means that
only error messages from gTest runs are printed; 'all' shows all logs.
'none' is used by default.

--gtest_repeat: it specifies the count of repeated runs. Use ten by default.

Any additional unknown args, such as --browser-ui-test-verify-pixels, are passed
to the gtest runner.
"""

import argparse
import hashlib
import pathlib
import shutil
import subprocess

# Constants used for color print.
_OK_GREEN = '\033[92m'
_FAIL_RED = '\033[91m'
_ENDC = '\033[0m'

# Used by the directory to host screenshots generated in each iteration. Add
# some special characters to make this name unique.
_TEMP_DIRECTORY_NAME_BASE = '@@check_pixel_test_flakiness!#'


class FlakyScreenshotError(Exception):
  """One of the screenshots has been detected to be flaky."""


class MissingScreenshotsError(Exception):
  """There were no screenshots found."""


def _get_md5(path):
  """Returns the Md5 digest of the specified file."""
  if not path.is_absolute():
    raise ValueError(f'{path} must be absolute')

  with path.open(mode='rb') as target_file:
    return hashlib.md5(target_file.read()).hexdigest()


def _compare_with_last_iteration(screenshots, prev_temp_dir, temp_dir,
                                 names_hash_mappings, flaky_screenshot_dir):
  """Compares the screenshots generated in the current iteration with those
  from the previous iteration. If flakiness is detected, returns a flaky
  screenshot's name. Otherwise, returns an empty string.

  Args:
    screenshots: A list of screenshot Paths.
    prev_temp_dir: The absolute file path to the directory that hosts the
      screenshots generated by the previous iteration.
    temp_dir: The absolute file path to the directory that hosts the screenshots
      generated by the current iteration.
    names_hash_mappings: The mappings from screenshot names to hash code.
    flaky_screenshot_dir: The absolute file path to the directory used to host
      flaky screenshots. If it is None, flaky screenshots are not written into
      files.

  Returns: None

  Raises:
    FlakyScreenshotError: if screenshots from prev_temp_dir do not match
      temp_dir
    """

  if prev_temp_dir is None:
    raise TypeError('prev_temp_dir is required to be a valid Path')

  if not screenshots:
    raise ValueError('screenshots must be non-empty')

  for screenshot in screenshots:
    # The screenshot hash code does not change so no flakiness is detected
    # on `screenshot`.
    if names_hash_mappings[screenshot.name] == _get_md5(screenshot):
      continue

    if flaky_screenshot_dir is not None:
      # Delete the output directory if it already exists.
      if flaky_screenshot_dir.exists():
        shutil.rmtree(flaky_screenshot_dir)

      flaky_screenshot_dir.mkdir(parents=True)

      # Move the screenshot generated by the last iteration to the dest
      # directory.
      shutil.move(
          prev_temp_dir / screenshot.name, flaky_screenshot_dir /
          f'{screenshot.stem}_Version_1{screenshot.suffix}')

      # Move the screenshot generated by the current iteration to the dest
      # directory.
      shutil.move(
          screenshot, flaky_screenshot_dir /
          f'{screenshot.stem}_Version_2{screenshot.suffix}')

    raise FlakyScreenshotError(
        f'{_FAIL_RED}[Failure]{_ENDC} Detect flakiness in: {screenshot.name}')

  # No flakiness detected.
  return None


def _analyze_screenshots(prev_temp_dir, temp_dir, names_hash_mappings,
                         flaky_screenshot_dir):
  """Analyzes the screenshots generated by one iteration.

  Args:
    prev_temp_dir: The absolute file path to the directory that hosts the
      screenshots generated by the previous iteration, or None if this is the
      first iteration.
    temp_dir: The absolute file path to the directory that hosts the screenshots
      generated by the current iteration.
    names_hash_mappings: The mappings from screenshot names to hash code.
    flaky_screenshot_dir: The absolute file path to the directory used to host
      flaky screenshots. If it is None, flaky screenshots are not written into
      files.

  Returns: None

  Raises:
    FlakyScreenshotError
    MissingScreenshotsError
  """
  screenshots = list(temp_dir.iterdir())
  if not screenshots:
    raise MissingScreenshotsError(
        f'{_FAIL_RED}[Failure]{_ENDC} no screenshots are generated in the '
        'specified tests: are you using the correct test filter?')

  # For the first iteration, nothing to compare with. Therefore, fill
  # `names_hash_mappings` and return.
  if prev_temp_dir is None:
    for screenshot in screenshots:
      names_hash_mappings[screenshot.name] = _get_md5(screenshot)
    return

  _compare_with_last_iteration(screenshots, prev_temp_dir, temp_dir,
                               names_hash_mappings, flaky_screenshot_dir)


def main():
  parser = argparse.ArgumentParser(
      description='Detect flakiness in the Skia Gold based pixel tests by '
      'running the specified pixel test executable file multiple iterations '
      'and comparing screenshots generated by neighboring iterations through '
      'file hash code. Warning: this script can only be used to detect '
      'flakiness in the pixel tests that use precise comparison.')
  parser.add_argument(
      '--test_target',
      type=str,
      required=True,
      help='a '
      'relative file path from the current working directory, or absolute file '
      'path, to the test executable based on Skia Gold, such as ash_pixeltests')
  parser.add_argument('--gtest_repeat',
                      type=int,
                      default=10,
                      help='the count of the repeated runs. The default value '
                      'is ten.')
  parser.add_argument('--root_dir',
                      type=str,
                      default='',
                      help='a relative file path from the current working '
                      'directory, or an absolute path, to the root directory '
                      'that hosts output data including the screenshots '
                      'generated in each iteration and the detected flaky '
                      'screenshots')
  parser.add_argument('--output_dir',
                      type=str,
                      help='a relative path starting from the output root path '
                      'specified by --root_dir or the current working '
                      'directory if --root_dir is omitted. It specifies a '
                      'directory used to host the flaky screenshots if any.')
  parser.add_argument('--log_mode',
                      choices=['none', 'error_only', 'all'],
                      default='none',
                      help='the option to control the log output during test '
                      'runs. `none` means that the log generated by test runs '
                      'does not show; `error_only` means that only error logs '
                      'are printed; `all` shows all logs. `none` is used by '
                      'default.')
  [known_args, unknown_args] = parser.parse_known_args()

  # Calculate the absolute path to the pixel test executable file.
  executable_full_path = pathlib.Path(known_args.test_target).resolve()

  # Leverage testing/xvfb.py to set up Aura env.
  xvfb_py_full_path = pathlib.Path('testing/xvfb.py').resolve()

  # Calculate the absolute path to the directory that hosts output data.
  output_root_path = pathlib.Path(known_args.root_dir).resolve()

  # Skip the Skia Gold functionality. Because this script compares images
  # through hash code.
  pixel_test_command_base = [
      str(xvfb_py_full_path),
      str(executable_full_path), '--bypass-skia-gold-functionality'
  ]

  # Pass unknown args to gtest.
  if unknown_args:
    pixel_test_command_base += unknown_args

  # Print the command to run pixel tests.
  full_command = ' '.join(pixel_test_command_base)
  print(f'{_OK_GREEN}[Begin]{_ENDC} {full_command}')

  # Configure log output.
  std_out_mode = subprocess.DEVNULL
  if known_args.log_mode == 'all':
    std_out_mode = None
  std_err_mode = None
  if known_args.log_mode == 'none':
    std_err_mode = subprocess.DEVNULL

  # Cache the screenshot host directory used in the last iteration. It updates
  # at the end of each iteration.
  prev_temp_dir = None

  # Similar to `prev_temp_dir` but it caches data for the active
  # iteration.
  temp_dir = None

  # Mappings screenshot names to hash code.
  names_hash_mappings = {}

  # Calculate the directory path for saving flaky screenshots.
  flaky_screenshot_dir = None
  if known_args.output_dir is not None:
    flaky_screenshot_dir = output_root_path / known_args.output_dir

  try:

    for i in range(known_args.gtest_repeat):
      # Calculate the absolute path to the screenshot host directory used for
      # this iteration. Recreate the host directory if it already exists.
      temp_dir = output_root_path / f'{_TEMP_DIRECTORY_NAME_BASE}{i}'

      if temp_dir.exists():
        shutil.rmtree(temp_dir)
      temp_dir.mkdir(parents=True)

      # Append the option so that the screenshots generated in pixel tests are
      # written into `temp_dir`.
      pixel_test_command = pixel_test_command_base[:]
      pixel_test_command.append(
          f'--skia-gold-local-png-write-directory={temp_dir}')

      # Run pixel tests.
      subprocess.run(pixel_test_command,
                     stdout=std_out_mode,
                     stderr=std_err_mode,
                     check=True)

      _analyze_screenshots(prev_temp_dir, temp_dir, names_hash_mappings,
                           flaky_screenshot_dir)

      print(f'{_OK_GREEN}[OK]{_ENDC} the iteration {i} succeeds')

      # Delete the temporary data directory used by the previous loop iteration
      # before overwriting it.
      if prev_temp_dir is not None:
        shutil.rmtree(prev_temp_dir)

      prev_temp_dir = temp_dir
    else:
      # The for loop has finished without exceptions.
      print(f'{_OK_GREEN}[Success]{_ENDC} no flakiness is detected')

  finally:
    # ensure that temp data are removed.
    for dir_to_rm in (prev_temp_dir, temp_dir):
      if dir_to_rm is not None and dir_to_rm.exists():
        shutil.rmtree(dir_to_rm)


if __name__ == '__main__':
  main()