#!/usr/bin/env vpython3
# Copyright 2020 The Chromium Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import json
import os
import shutil
import sys
import tempfile
from textwrap import dedent
import unittest

# The following non-std imports are fetched via vpython. See the list at
# //.vpython3
import mock  # pylint: disable=import-error
from parameterized import parameterized  # pylint: disable=import-error

import test_runner

_TAST_TEST_RESULTS_JSON = {
    "name": "login.Chrome",
    "errors": None,
    "start": "2020-01-01T15:41:30.799228462-08:00",
    "end": "2020-01-01T15:41:53.318914698-08:00",
    "skipReason": ""
}


class TestRunnerTest(unittest.TestCase):

  def setUp(self):
    self._tmp_dir = tempfile.mkdtemp()
    self.mock_rdb = mock.patch.object(
        test_runner.result_sink, 'TryInitClient', return_value=None)
    self.mock_rdb.start()
    self.mock_env = mock.patch.dict(
        os.environ, {'SWARMING_BOT_ID': 'cros-chrome-chromeos8-row29'})
    self.mock_env.start()

  def tearDown(self):
    shutil.rmtree(self._tmp_dir, ignore_errors=True)
    self.mock_rdb.stop()
    self.mock_env.stop()

  def safeAssertItemsEqual(self, list1, list2):
    """A Py3 safe version of assertItemsEqual.

    See https://bugs.python.org/issue17866.
    """
    self.assertSetEqual(set(list1), set(list2))


class TastTests(TestRunnerTest):

  def get_common_tast_args(self, use_vm, fetch_cros_hostname):
    return [
        'script_name',
        'tast',
        '--suite-name=chrome_all_tast_tests',
        '--board=eve',
        '--flash',
        '--path-to-outdir=out_eve/Release',
        '--logs-dir=%s' % self._tmp_dir,
        '--use-vm' if use_vm else
        ('--fetch-cros-hostname'
         if fetch_cros_hostname else '--device=localhost:2222'),
    ]

  def get_common_tast_expectations(self, use_vm, fetch_cros_hostname):
    expectation = [
        test_runner.CROS_RUN_TEST_PATH,
        '--board',
        'eve',
        '--cache-dir',
        test_runner.DEFAULT_CROS_CACHE,
        '--results-dest-dir',
        '%s/system_logs' % self._tmp_dir,
        '--flash',
        '--build-dir',
        'out_eve/Release',
        '--results-dir',
        self._tmp_dir,
        '--tast-total-shards=1',
        '--tast-shard-index=0',
    ]
    expectation.extend(['--start', '--copy-on-write'] if use_vm else (
        ['--device', 'chrome-chromeos8-row29']
        if fetch_cros_hostname else ['--device', 'localhost:2222']))
    for p in test_runner.SYSTEM_LOG_LOCATIONS:
      expectation.extend(['--results-src', p])

    expectation += [
        '--mount',
        '--deploy',
        '--nostrip',
    ]
    return expectation

  def test_tast_gtest_filter(self):
    """Tests running tast tests with a gtest-style filter."""
    with open(os.path.join(self._tmp_dir, 'streamed_results.jsonl'), 'w') as f:
      json.dump(_TAST_TEST_RESULTS_JSON, f)

    args = self.get_common_tast_args(False, False) + [
        '--attr-expr=( "group:mainline" && "dep:chrome" && !informational)',
        '--gtest_filter=login.Chrome:ui.WindowControl',
    ]
    with mock.patch.object(sys, 'argv', args),\
         mock.patch.object(test_runner.subprocess, 'Popen') as mock_popen:
      mock_popen.return_value.returncode = 0

      test_runner.main()
      # The gtest filter should cause the Tast expr to be replaced with a list
      # of the tests in the filter.
      expected_cmd = self.get_common_tast_expectations(False, False) + [
          '--tast=("name:login.Chrome" || "name:ui.WindowControl")'
      ]

      self.safeAssertItemsEqual(expected_cmd, mock_popen.call_args[0][0])

  @parameterized.expand([
      [True, False],
      [False, True],
      [False, False],
  ])
  def test_tast_attr_expr(self, use_vm, fetch_cros_hostname):
    """Tests running a tast tests specified by an attribute expression."""
    with open(os.path.join(self._tmp_dir, 'streamed_results.jsonl'), 'w') as f:
      json.dump(_TAST_TEST_RESULTS_JSON, f)

    args = self.get_common_tast_args(use_vm, fetch_cros_hostname) + [
        '--attr-expr=( "group:mainline" && "dep:chrome" && !informational)',
    ]
    with mock.patch.object(sys, 'argv', args),\
         mock.patch.object(test_runner.subprocess, 'Popen') as mock_popen:
      mock_popen.return_value.returncode = 0

      test_runner.main()
      expected_cmd = self.get_common_tast_expectations(
          use_vm, fetch_cros_hostname) + [
              '--tast=( "group:mainline" && "dep:chrome" && !informational)',
          ]

      self.safeAssertItemsEqual(expected_cmd, mock_popen.call_args[0][0])

  @parameterized.expand([
      [True, False],
      [False, True],
      [False, False],
  ])
  def test_tast_with_vars(self, use_vm, fetch_cros_hostname):
    """Tests running a tast tests with runtime variables."""
    with open(os.path.join(self._tmp_dir, 'streamed_results.jsonl'), 'w') as f:
      json.dump(_TAST_TEST_RESULTS_JSON, f)

    args = self.get_common_tast_args(use_vm, fetch_cros_hostname) + [
        '-t=login.Chrome',
        '--tast-var=key=value',
    ]
    with mock.patch.object(sys, 'argv', args),\
         mock.patch.object(test_runner.subprocess, 'Popen') as mock_popen:
      mock_popen.return_value.returncode = 0
      test_runner.main()
      expected_cmd = self.get_common_tast_expectations(
          use_vm, fetch_cros_hostname) + [
              '--tast', 'login.Chrome', '--tast-var', 'key=value'
          ]

      self.safeAssertItemsEqual(expected_cmd, mock_popen.call_args[0][0])

  @parameterized.expand([
      [True, False],
      [False, True],
      [False, False],
  ])
  def test_tast_retries(self, use_vm, fetch_cros_hostname):
    """Tests running a tast tests with retries."""
    with open(os.path.join(self._tmp_dir, 'streamed_results.jsonl'), 'w') as f:
      json.dump(_TAST_TEST_RESULTS_JSON, f)

    args = self.get_common_tast_args(use_vm, fetch_cros_hostname) + [
        '-t=login.Chrome',
        '--tast-retries=1',
    ]
    with mock.patch.object(sys, 'argv', args),\
         mock.patch.object(test_runner.subprocess, 'Popen') as mock_popen:
      mock_popen.return_value.returncode = 0
      test_runner.main()
      expected_cmd = self.get_common_tast_expectations(
          use_vm,
          fetch_cros_hostname) + ['--tast', 'login.Chrome', '--tast-retries=1']

      self.safeAssertItemsEqual(expected_cmd, mock_popen.call_args[0][0])

  @parameterized.expand([
      [True, False],
      [False, True],
      [False, False],
  ])
  def test_tast(self, use_vm, fetch_cros_hostname):
    """Tests running a tast tests."""
    with open(os.path.join(self._tmp_dir, 'streamed_results.jsonl'), 'w') as f:
      json.dump(_TAST_TEST_RESULTS_JSON, f)

    args = self.get_common_tast_args(use_vm, fetch_cros_hostname) + [
        '-t=login.Chrome',
    ]
    with mock.patch.object(sys, 'argv', args),\
         mock.patch.object(test_runner.subprocess, 'Popen') as mock_popen:
      mock_popen.return_value.returncode = 0

      test_runner.main()
      expected_cmd = self.get_common_tast_expectations(
          use_vm, fetch_cros_hostname) + ['--tast', 'login.Chrome']

      self.safeAssertItemsEqual(expected_cmd, mock_popen.call_args[0][0])


class GTestTest(TestRunnerTest):

  @parameterized.expand([
      [True, True, True, False, True],
      [True, False, False, False, False],
      [False, True, True, True, True],
      [False, False, False, True, False],
      [False, True, True, False, True],
      [False, False, False, False, False],
  ])
  def test_gtest(self, use_vm, stop_ui, use_test_sudo_helper,
                 fetch_cros_hostname, use_deployed_dbus_configs):
    """Tests running a gtest."""
    fd_mock = mock.mock_open()

    args = [
        'script_name',
        'gtest',
        '--test-exe=out_eve/Release/base_unittests',
        '--board=eve',
        '--path-to-outdir=out_eve/Release',
        '--use-vm' if use_vm else
        ('--fetch-cros-hostname'
         if fetch_cros_hostname else '--device=localhost:2222'),
    ]
    if stop_ui:
      args.append('--stop-ui')
    if use_test_sudo_helper:
      args.append('--run-test-sudo-helper')
    if use_deployed_dbus_configs:
      args.append('--use-deployed-dbus-configs')

    with mock.patch.object(sys, 'argv', args),\
         mock.patch.object(test_runner.subprocess, 'Popen') as mock_popen,\
         mock.patch.object(os, 'fdopen', fd_mock),\
         mock.patch.object(os, 'remove') as mock_remove,\
         mock.patch.object(tempfile, 'mkstemp',
            side_effect=[(3, 'out_eve/Release/device_script.sh'),\
                         (4, 'out_eve/Release/runtime_files.txt')]),\
         mock.patch.object(os, 'fchmod'):
      mock_popen.return_value.returncode = 0

      test_runner.main()
      self.assertEqual(1, mock_popen.call_count)
      expected_cmd = [
          'vpython3', test_runner.CROS_RUN_TEST_PATH, '--board', 'eve',
          '--cache-dir', test_runner.DEFAULT_CROS_CACHE, '--remote-cmd',
          '--cwd', 'out_eve/Release', '--files-from',
          'out_eve/Release/runtime_files.txt'
      ]
      if not stop_ui:
        expected_cmd.append('--as-chronos')
      expected_cmd.extend(['--start', '--copy-on-write'] if use_vm else (
          ['--device', 'chrome-chromeos8-row29']
          if fetch_cros_hostname else ['--device', 'localhost:2222']))
      expected_cmd.extend(['--', './device_script.sh'])
      self.safeAssertItemsEqual(expected_cmd, mock_popen.call_args[0][0])

      expected_device_script = dedent("""\
          #!/bin/sh
          export HOME=/usr/local/tmp
          export TMPDIR=/usr/local/tmp
          """)

      core_cmd = 'LD_LIBRARY_PATH=./ ./out_eve/Release/base_unittests'\
          ' --test-launcher-shard-index=0 --test-launcher-total-shards=1'

      if use_test_sudo_helper:
        expected_device_script += dedent("""\
            TEST_SUDO_HELPER_PATH=$(mktemp)
            ./test_sudo_helper.py --socket-path=${TEST_SUDO_HELPER_PATH} &
            TEST_SUDO_HELPER_PID=$!
          """)
        core_cmd += ' --test-sudo-helper-socket-path=${TEST_SUDO_HELPER_PATH}'

      if use_deployed_dbus_configs:
        expected_device_script += dedent("""\
            mount --bind ./dbus /opt/google/chrome/dbus
            kill -s HUP $(pgrep dbus)
          """)

      if stop_ui:
        dbus_cmd = 'dbus-send --system --type=method_call'\
          ' --dest=org.chromium.PowerManager'\
          ' /org/chromium/PowerManager'\
          ' org.chromium.PowerManager.HandleUserActivity int32:0'
        expected_device_script += dedent("""\
          stop ui
          {0}
          chown -R chronos: ../..
          sudo -E -u chronos -- /bin/bash -c \"{1}\"
          TEST_RETURN_CODE=$?
          start ui
          """).format(dbus_cmd, core_cmd)
      else:
        expected_device_script += dedent("""\
          {0}
          TEST_RETURN_CODE=$?
          """).format(core_cmd)

      if use_test_sudo_helper:
        expected_device_script += dedent("""\
            pkill -P $TEST_SUDO_HELPER_PID
            kill $TEST_SUDO_HELPER_PID
            unlink ${TEST_SUDO_HELPER_PATH}
          """)

      if use_deployed_dbus_configs:
        expected_device_script += dedent("""\
            umount /opt/google/chrome/dbus
            kill -s HUP $(pgrep dbus)
          """)

      expected_device_script += dedent("""\
          exit $TEST_RETURN_CODE
        """)

      self.assertEqual(2, fd_mock().write.call_count)
      write_calls = fd_mock().write.call_args_list

      # Split the strings to make failure messages easier to read.
      # Verify the first write of device script.
      self.assertListEqual(
          expected_device_script.split('\n'),
          str(write_calls[0][0][0]).split('\n'))

      # Verify the 2nd write of runtime files.
      expected_runtime_files = ['out_eve/Release/device_script.sh']
      self.assertListEqual(expected_runtime_files,
                           str(write_calls[1][0][0]).strip().split('\n'))

      mock_remove.assert_called_once_with('out_eve/Release/device_script.sh')

  def test_gtest_with_vpython(self):
    """Tests building a gtest with --vpython-dir."""
    args = mock.MagicMock()
    args.test_exe = 'base_unittests'
    args.test_launcher_summary_output = None
    args.trace_dir = None
    args.runtime_deps_path = None
    args.path_to_outdir = self._tmp_dir
    args.vpython_dir = self._tmp_dir
    args.logs_dir = self._tmp_dir

    # With vpython_dir initially empty, the test_runner should error out
    # due to missing vpython binaries.
    gtest = test_runner.GTestTest(args, None)
    with self.assertRaises(test_runner.TestFormatError):
      gtest.build_test_command()

    # Create the two expected tools, and the test should be ready to run.
    with open(os.path.join(args.vpython_dir, 'vpython3'), 'w'):
      pass  # Just touch the file.
    os.mkdir(os.path.join(args.vpython_dir, 'bin'))
    with open(os.path.join(args.vpython_dir, 'bin', 'python3'), 'w'):
      pass
    gtest = test_runner.GTestTest(args, None)
    gtest.build_test_command()


class HostCmdTests(TestRunnerTest):

  @parameterized.expand([
      [False, False],
      [False, True],
      [True, False],
      [True, True],
  ])
  def test_host_cmd(self, deploy_chrome, strip_chrome):
    args = [
        'script_name',
        'host-cmd',
        '--board=eve',
        '--flash',
        '--path-to-outdir=out/Release',
        '--device=localhost:2222',
    ]
    if deploy_chrome:
      args += ['--deploy-chrome']
    if strip_chrome:
      args += ['--strip-chrome']
    args += [
        '--',
        'fake_cmd',
    ]
    with mock.patch.object(sys, 'argv', args),\
         mock.patch.object(test_runner.subprocess, 'Popen') as mock_popen:
      mock_popen.return_value.returncode = 0

      test_runner.main()
      expected_cmd = [
          test_runner.CROS_RUN_TEST_PATH,
          '--board',
          'eve',
          '--cache-dir',
          test_runner.DEFAULT_CROS_CACHE,
          '--flash',
          '--device',
          'localhost:2222',
          '--host-cmd',
      ]
      if deploy_chrome:
        expected_cmd += [
            '--mount',
            '--deploy',
            '--build-dir',
            os.path.join(test_runner.CHROMIUM_SRC_PATH, 'out/Release'),
        ]
        if not strip_chrome:
          expected_cmd += ['--nostrip']

      expected_cmd += [
          '--',
          'fake_cmd',
      ]

      self.safeAssertItemsEqual(expected_cmd, mock_popen.call_args[0][0])


if __name__ == '__main__':
  unittest.main()