# 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.

import contextlib
import json
import pathlib
import unittest
from typing import TYPE_CHECKING
from unittest import mock

from typing_extensions import override

from crossbench.browsers.browser import Browser
from crossbench.browsers.webdriver import RemoteWebDriver
from crossbench.env.runner_env import RunnerEnv
from crossbench.exception import MultiException
from crossbench.flags.base import Flags
from crossbench.helper.state import UnexpectedStateError
from crossbench.probes import all as all_probes
from crossbench.probes.probe import ProbeIncompatibleBrowser
from crossbench.runner.groups.session import BrowserSessionRunGroup
from crossbench.runner.groups.thread import RunThreadGroup
from crossbench.runner.runner import Runner, ThreadMode
from tests import test_helper
from tests.crossbench.mock_browser import MockChromeDev
from tests.crossbench.mock_helper import MockBenchmark
from tests.crossbench.runner.helper import (BaseRunnerTestCase, MockBrowser,
                                            MockPlatform, MockProbe,
                                            MockProbeContext, MockRun,
                                            MockRunner)

if TYPE_CHECKING:
  from crossbench.probes.probe import Probe

# Skip strict type checks for better mocking
# pytype: disable=wrong-arg-types
class TestThreadModeTestCase(unittest.TestCase):
  # pylint has some issues with enums.
  # pylint: disable=no-member

  def create_session(self, browser, index) -> BrowserSessionRunGroup:
    return BrowserSessionRunGroup(
        self.env,
        self.probes,
        browser,
        Flags(),
        index,
        self.root_dir,
        create_symlinks=True,
        throw=True)

  @override
  def setUp(self) -> None:
    self.platform_a = MockPlatform("platform a")
    self.platform_b = MockPlatform("platform b")
    self.browser_a_1 = MockBrowser("mock browser a 1", self.platform_a)
    self.browser_a_2 = MockBrowser("mock browser b 1", self.platform_a)
    self.browser_b_1 = MockBrowser("mock browser b 1", self.platform_b)
    self.browser_b_2 = MockBrowser("mock browser b 2", self.platform_b)
    self.runner = MockRunner()
    self.root_dir = pathlib.Path()
    self.env = self.runner.env
    self.probes: list[Probe] = []
    self.runs = (
        MockRun(self.runner, self.create_session(self.browser_a_1, 1), "run 1"),
        MockRun(self.runner, self.create_session(self.browser_a_2, 2), "run 2"),
        MockRun(self.runner, self.create_session(self.browser_a_1, 3), "run 3"),
        MockRun(self.runner, self.create_session(self.browser_a_2, 4), "run 4"),
        MockRun(self.runner, self.create_session(self.browser_b_1, 5), "run 5"),
        MockRun(self.runner, self.create_session(self.browser_b_2, 6), "run 6"),
        MockRun(self.runner, self.create_session(self.browser_b_1, 7), "run 7"),
        MockRun(self.runner, self.create_session(self.browser_b_2, 8), "run 8"),
    )
    self.runner.runs = self.runs

  def test_default_runs(self):
    session_ids = {run.browser_session.index for run in self.runs}
    self.assertEqual(len(session_ids), len(self.runs))

  def test_group_none(self):
    groups = ThreadMode.NONE.group(self.runs)
    self.assertEqual(len(groups), 1)
    self.assertTupleEqual(groups[0].runs, self.runs)
    self.assertEqual(groups[0].index, 0)

  def test_group_platform(self):
    groups = ThreadMode.PLATFORM.group(self.runs)
    self.assertEqual(len(groups), 2)
    group_a, group_b = groups
    self.assertTupleEqual(group_a.runs, self.runs[:4])
    self.assertTupleEqual(group_b.runs, self.runs[4:])
    self.assertEqual(group_a.index, 0)
    self.assertEqual(group_b.index, 1)

  def test_group_browser(self):
    groups = ThreadMode.BROWSER.group(self.runs)
    self.assertEqual(len(groups), 4)
    self.assertTupleEqual(groups[0].runs, (self.runs[0], self.runs[2]))
    self.assertTupleEqual(groups[1].runs, (self.runs[1], self.runs[3]))
    self.assertTupleEqual(groups[2].runs, (self.runs[4], self.runs[6]))
    self.assertTupleEqual(groups[3].runs, (self.runs[5], self.runs[7]))
    for index, group in enumerate(groups):
      self.assertEqual(group.index, index)

  def test_group_session(self):
    groups = ThreadMode.SESSION.group(self.runs)
    self.assertEqual(len(groups), len(self.runs))
    for group, run in zip(groups, self.runs):
      self.assertTupleEqual(group.runs, (run,))
    for index, group in enumerate(groups):
      self.assertEqual(group.index, index)

  def test_group_session_2(self):
    session_1 = self.create_session(self.browser_a_1, 1)
    session_2 = self.create_session(self.browser_a_2, 2)
    runs = (
        MockRun(self.runner, session_1, "story 1"),
        MockRun(self.runner, session_2, "story 2"),
        MockRun(self.runner, session_1, "story 3"),
        MockRun(self.runner, session_2, "story 4"),
    )
    groups = ThreadMode.SESSION.group(runs)
    group_a, group_b = groups
    self.assertTupleEqual(group_a.runs, (runs[0], runs[2]))
    self.assertTupleEqual(group_b.runs, (runs[1], runs[3]))
    for index, group in enumerate(groups):
      self.assertEqual(group.index, index)


class FailingMockProbeContext(MockProbeContext):

  @override
  def setup(self):
    raise CustomException("failing setup")


class RunnerTestCase(BaseRunnerTestCase):

  def test_default_instance(self):
    runner = self.default_runner()
    self.assertSequenceEqual(self.stories, runner.stories)
    self.assertSequenceEqual(self.browsers, runner.browsers)
    self.assertEqual(runner.repetitions, 1)
    self.assertEqual(len(runner.platforms), 1)
    self.assertTrue(runner.exceptions.is_success)
    default_probes = list(runner.default_probes)
    self.assertListEqual(list(runner.probes), default_probes)
    self.assertEqual(
        len(default_probes), len(all_probes.DEFAULT_INTERNAL_PROBES))
    self.assertEqual(len(runner.runs), 0)
    # no runs => is_success == false
    self.assertFalse(runner.is_success)

  def test_dry_run(self):
    self.test_run(is_dry_run=True)

  def test_run(self, is_dry_run=False):
    runner = self.default_runner()

    runner.run(is_dry_run)
    # Don't reuse the Runner:
    with self.assertRaises(UnexpectedStateError):
      runner.run(is_dry_run)

    self.assertEqual(len(runner.runs), 4)
    self.assertTrue(runner.is_success)
    for run in runner.runs:
      self.assertTrue(run.is_success)

      if not is_dry_run:
        self.assertEqual(
            len(run.results), len(all_probes.DEFAULT_INTERNAL_PROBES))
        for probe in runner.probes:
          self.assertIn(probe, run.results)

  def test_run_mock_probe(self):
    runner = self.default_runner()
    probe = MockProbe("custom_probe_data")
    runner.attach_probe(probe)
    self.assertIn(probe, runner.probes)
    for browser in runner.browsers:
      self.assertIn(probe, browser.probes)

    runner.run()
    self.assertTrue(runner.is_success)
    self.assertEqual(len(runner.runs), 4)
    for run in runner.runs:
      self._validate_successful_run(run, runner, probe, "custom_probe_data")
    for browser in runner.browsers:
      runs_symlinks = list(
          (runner.out_dir / browser.unique_name / "runs").iterdir())
      self.assertEqual(len(runs_symlinks), 2)

  def test_run_remote_web_driver(self):
    driver = mock.Mock()
    driver.capabilities = {
        "browserVersion": "123.0.4567.89",
        "setWindowRect": False,
    }
    browser = RemoteWebDriver("test-driver", driver)
    runner = self.default_runner(browsers=[browser])
    runner.run()

  def _validate_successful_run(self, run, runner, probe, probe_data):
    results = run.results[probe]
    with results.json.open() as f:
      probe_data = json.load(f)
      self.assertEqual(probe_data, probe_data)
    browser_dir = runner.out_dir / run.browser.unique_name
    # Pyfakefs is having some issues with relative symlinks, thus we're
    # manually combining the paths.
    runs_dir = browser_dir / "runs"
    run_symlink = runs_dir / (runs_dir / str(run.index)).readlink()
    self.assertEqual(run_symlink.resolve(), run.out_dir)
    self._validate_internal_probes(run, runner)

  def _validate_internal_probes(self, run, runner):
    for probe in runner.probes:
      if not probe.is_internal:
        continue
      result = run.results[probe]
      self.assertTrue(result)

  def test_single_story_run_mock_probe_partial_setup_fail(self):
    runner = self.single_story_runner(throw=False)

    probe = MockProbe("custom_probe_data", FailingMockProbeContext)
    runner.attach_probe(probe)
    self.assertIn(probe, runner.probes)
    for browser in runner.browsers:
      self.assertIn(probe, browser.probes)

    with self.assertRaises(MultiException) as cm:
      runner.run()
    self.assertEqual(len(cm.exception), 1)
    exception = cm.exception.exceptions[0].exception
    self.assertIsInstance(exception, CustomException)

    self.assertFalse(runner.is_success)
    self.assertEqual(len(runner.runs), 1)
    failed_run = list(runner.runs)[0]
    self.assertFalse(failed_run.is_success)
    self.assertTrue(failed_run.results[probe].is_empty)
    self._validate_internal_probes(failed_run, runner)

  def test_single_story_run_mock_probe_calls(self):
    # Make sure start / stop are called.
    runner = self.single_story_runner(throw=True)
    with mock.patch.object(FailingMockProbeContext,
                           "setup") as setup_mock, mock.patch.object(
                               FailingMockProbeContext,
                               "start") as start_mock, mock.patch.object(
                                   FailingMockProbeContext,
                                   "stop") as stop_mock:
      probe = MockProbe("custom_probe_data", FailingMockProbeContext)
      runner.attach_probe(probe)
      runner.run()
    self.assertTrue(runner.is_success)
    setup_mock.assert_called_once()
    start_mock.assert_called_once()
    stop_mock.assert_called_once()

  def test_single_story_run_mock_probe_partial_setup_fail_mock(self):
    # Make sure start / stop / teardown are not called after a setup failure
    runner = self.single_story_runner(throw=False)
    with mock.patch.object(FailingMockProbeContext,
                           "start") as start_mock, mock.patch.object(
                               FailingMockProbeContext,
                               "stop") as stop_mock, mock.patch.object(
                                   FailingMockProbeContext,
                                   "teardown") as teardown_mock:
      probe = MockProbe("custom_probe_data", FailingMockProbeContext)
      runner.attach_probe(probe)
      with self.assertRaises(MultiException) as cm:
        runner.run()
      exception = cm.exception.exceptions[0].exception
      self.assertIsInstance(exception, CustomException)
    self.assertFalse(runner.is_success)
    start_mock.assert_not_called()
    stop_mock.assert_not_called()
    teardown_mock.assert_not_called()

  def test_run_mock_probe_partial_setup_fail(self):
    runner = self.default_runner(throw=False)
    setup_count = 0

    class PartialFailingMockProbeContext(MockProbeContext):

      @override
      def setup(self):
        nonlocal setup_count
        setup_count += 1
        if setup_count == 3:
          raise CustomException(f"failing setup number {setup_count}")

    probe = MockProbe("custom_probe_data", PartialFailingMockProbeContext)
    runner.attach_probe(probe)
    self.assertIn(probe, runner.probes)
    for browser in runner.browsers:
      self.assertIn(probe, browser.probes)

    with self.assertRaises(MultiException) as cm:
      runner.run()
    self.assertEqual(len(cm.exception), 1)
    exception = cm.exception.exceptions[0].exception
    self.assertIsInstance(exception, CustomException)

    self.assertFalse(runner.is_success)
    self.assertEqual(setup_count, 4)
    self.assertEqual(len(runner.runs), 4)
    failed_runs = list(run for run in runner.runs if not run.is_success)
    self.assertEqual(len(failed_runs), 1)
    failed_run = failed_runs[0]

    for run in runner.runs:
      if run is failed_run:
        continue
      self.assertTrue(run.is_success)
      self._validate_successful_run(run, runner, probe, "custom_probe_data")

    self.assertEqual(failed_run.index, 2)
    self.assertFalse(failed_run.is_success)
    self.assertTrue(failed_run.results[probe].is_empty)
    self._validate_internal_probes(failed_run, runner)

  def test_attach_probe_twice(self):
    runner = self.default_runner()
    probe = MockProbe("custom_probe_data")
    runner.attach_probe(probe)
    # Cannot attach same probe twice.
    with self.assertRaises(ValueError) as cm:
      runner.attach_probe(probe)
    self.assertIn("twice", str(cm.exception))
    self.assertIn(probe, runner.probes)
    self.assertNotIn(probe, runner.default_probes)

  def test_attach_incompatible_probe(self):
    runner = self.default_runner()
    probe = MockProbe("custom_probe_data")

    def mock_validate_browser(env: RunnerEnv, browser: Browser):
      del env
      nonlocal probe
      raise ProbeIncompatibleBrowser(probe, browser, "mock invalid")

    probe.validate_browser = mock_validate_browser
    with self.assertRaises(ProbeIncompatibleBrowser) as cm:
      runner.attach_probe(probe)
    self.assertIn("mock invalid", str(cm.exception))
    # matching_browser_only = True silence the error
    runner.attach_probe(probe, matching_browser_only=True)
    # No browser matches => probe is not available
    self.assertNotIn(probe, runner.probes)
    self.assertNotIn(probe, runner.default_probes)
    for browser in self.browsers:
      self.assertNotIn(probe, browser.probes)

  def test_attach_partially_incompatible_probe(self):
    runner = self.default_runner()
    probe = MockProbe("custom_probe_data")
    compatible_browser = self.browsers[1]

    def mock_validate_browser(env: RunnerEnv, browser: Browser):
      del env
      nonlocal probe
      nonlocal compatible_browser
      if browser != compatible_browser:
        raise ProbeIncompatibleBrowser(probe, browser, "mock invalid")

    # Attaching incompatible probes raises errors by default.
    probe.validate_browser = mock_validate_browser
    with self.assertRaises(ProbeIncompatibleBrowser) as cm:
      runner.attach_probe(probe)
    self.assertIn("mock invalid", str(cm.exception))
    # matching_browser_only = True silences the error
    runner.attach_probe(probe, matching_browser_only=True)
    self.assertIn(probe, runner.probes)
    self.assertNotIn(probe, runner.default_probes)
    for browser in self.browsers:
      if browser == compatible_browser:
        self.assertIn(probe, browser.probes)
      else:
        self.assertNotIn(probe, browser.probes)


class CustomException(Exception):
  pass


class RunThreadGroupTestCase(BaseRunnerTestCase):

  @override
  def tearDown(self) -> None:
    for browser in self.browsers:
      self.assertFalse(browser.is_running)
    return super().tearDown()

  def test_create_no_runs(self):
    with self.assertRaises(AssertionError):
      RunThreadGroup([])

  def test_different_runners(self):
    runs_a = list(self.default_runner().get_runs())
    self.out_dir = self.out_dir.parent / "second_out_dir"
    runner_b = Runner(
        self.out_dir, [MockChromeDev("chrome-dev-2")],
        self.benchmark,
        platform=self.platform,
        throw=True,
        in_memory_result_db=True)
    runs_b = list(runner_b.get_runs())
    self.assertNotEqual(runs_a[0].runner, runs_b[0].runner)
    with self.assertRaises(AssertionError) as cm:
      RunThreadGroup(runs_a + runs_b)
    self.assertIn("same Runner", str(cm.exception))

  @contextlib.contextmanager
  def patch_teardown_run(self, runner):
    with mock.patch.object(
        runner.results_db, "teardown_run",
        side_effect=None) as teardown_run_mock:
      yield teardown_run_mock


  def test_simple_runs(self):
    runner = self.default_runner()
    runs = tuple(runner.get_runs())
    thread = RunThreadGroup(runs)
    self.assertEqual(thread.index, 0)
    self.assertEqual(thread.runner, runner)
    self.assertSequenceEqual(thread.runs, runs)
    self.assertTrue(thread.is_success)

    run_count = 0

    def test_run(run_method):
      nonlocal run_count
      run_count += 1
      run_method(is_dry_run=False)

    for run in runs:
      run.run = (  # pylint: disable=unnecessary-direct-lambda-call
          lambda run_method: lambda is_dry_run: test_run(run_method))(
              run.run)
    with self.patch_teardown_run(runner) as teardown_run_mock:
      thread.run()
      self.assertEqual(teardown_run_mock.call_count, len(runs))

    self.assertTrue(thread.is_success)
    self.assertSequenceEqual(thread.runs, runs)
    self.assertEqual(run_count, 4)

  def test_run_fail_run_probe_create_context(self):
    # 2 runs, same browser different stories
    runner = self.default_runner(browsers=[self.browsers[1]], throw=False)
    probe = MockProbe("custom_probe_data")
    runner.attach_probe(probe)
    self.assertTrue(probe.is_attached)
    runs = tuple(runner.get_runs())
    thread = RunThreadGroup(runs)
    failing_session, successful_session = thread.browser_sessions
    failing_run, successful_run = runs

    setup_fail_count = 0

    def mock_get_context_fail(run):
      if run == successful_run:
        return MockProbeContext(probe, run)
      nonlocal setup_fail_count
      setup_fail_count += 1
      raise CustomException()

    probe.create_context = mock_get_context_fail

    self.assertEqual(setup_fail_count, 0)
    with self.patch_teardown_run(runner) as teardown_run_mock:
      thread.run()
      self.assertEqual(teardown_run_mock.call_count, len(runs))
    self.assertEqual(setup_fail_count, 1)

    self.assertTrue(successful_session.is_success)
    self.assertTrue(successful_run.is_success)

    # Errors are propagated up:
    for exceptions_holder in (runner, thread, failing_session, failing_run):
      self.assertFalse(exceptions_holder.is_success)
      exceptions = exceptions_holder.exceptions
      self.assertEqual(len(exceptions), 1)
      exception_entry = exceptions[0]
      self.assertIsInstance(exception_entry.exception, CustomException)

  def test_run_fail_run_probe_setup(self):
    # 2 runs, same browser different stories
    runner = self.default_runner(browsers=[self.browsers[1]], throw=False)
    probe = MockProbe("custom_probe_data")
    runner.attach_probe(probe)
    self.assertTrue(probe.is_attached)
    runs = tuple(runner.get_runs())
    thread = RunThreadGroup(runs)
    failing_session, successful_session = thread.browser_sessions
    failing_run, successful_run = runs

    setup_fail_count = 0

    def mock_setup_fail() -> None:
      nonlocal setup_fail_count
      setup_fail_count += 1
      raise CustomException()

    def mock_get_context_fail(run):
      context = MockProbeContext(probe, run)
      if run == failing_run:
        context.setup = mock_setup_fail
      return context

    probe.create_context = mock_get_context_fail

    self.assertEqual(setup_fail_count, 0)
    with self.patch_teardown_run(runner) as teardown_run_mock:
      thread.run()
      self.assertEqual(teardown_run_mock.call_count, len(runs))
    self.assertEqual(setup_fail_count, 1)

    self.assertTrue(successful_session.is_success)
    self.assertTrue(successful_run.is_success)

    # Errors are propagated up:
    for exceptions_holder in (runner, thread, failing_session, failing_run):
      self.assertFalse(exceptions_holder.is_success)
      exceptions = exceptions_holder.exceptions
      self.assertEqual(len(exceptions), 1)
      exception_entry = exceptions[0]
      self.assertIsInstance(exception_entry.exception, CustomException)

  def test_run_fail_one_browser_setup(self):
    # 2 runs, same story, different browsers
    benchmark = MockBenchmark(stories=[self.stories[0]])
    runner = Runner(
        self.out_dir,
        self.browsers,
        benchmark,
        platform=self.platform,
        in_memory_result_db=True)
    runs = tuple(runner.get_runs())
    thread = RunThreadGroup(runs)
    failing_session, successful_session = thread.browser_sessions
    failing_run, successful_run = runs
    self.assertNotEqual(failing_run.browser, successful_run.browser)

    setup_fail_count = 0

    def mock_start_fail(session: BrowserSessionRunGroup) -> None:
      del session
      nonlocal setup_fail_count
      setup_fail_count += 1
      raise CustomException()

    failing_run.browser.start = mock_start_fail

    self.assertEqual(setup_fail_count, 0)
    with self.patch_teardown_run(runner) as teardown_run_mock:
      thread.run()
      self.assertEqual(teardown_run_mock.call_count, len(runs))
    self.assertEqual(setup_fail_count, 1)

    self.assertTrue(successful_session.is_success)
    self.assertTrue(successful_run.is_success)

    # browser startup failures should also propagate down to all runs.
    for exceptions_holder in (runner, thread, failing_session, failing_run):
      self.assertFalse(exceptions_holder.is_success)
      exceptions = exceptions_holder.exceptions
      self.assertEqual(len(exceptions), 1)
      exception_entry = exceptions[0]
      self.assertIsInstance(exception_entry.exception, CustomException)

  def test_run_fail_run(self):
    # 4 runs = (2 browser) x (2 stories)
    runner = self.default_runner(throw=False)
    runs = tuple(runner.get_runs())
    thread = RunThreadGroup(runs)
    failing_run = runs[0]
    failing_session = failing_run.browser_session

    run_fail_count = 0

    def mock_run_story_fail():
      nonlocal run_fail_count
      run_fail_count += 1
      raise CustomException()

    with mock.patch.object(failing_run, "_run_story", mock_run_story_fail):
      self.assertEqual(run_fail_count, 0)
      with self.patch_teardown_run(runner) as teardown_run_mock:
        thread.run()
        self.assertEqual(teardown_run_mock.call_count, len(runs))
      self.assertEqual(run_fail_count, 1)

    for session in thread.browser_sessions:
      if session != failing_run.browser_session:
        self.assertTrue(session.is_success)
    for run in runs:
      if run != failing_run:
        self.assertTrue(run.is_success)

    # Errors are propagate up:
    for exceptions_holder in (runner, thread, failing_session, failing_run):
      self.assertFalse(exceptions_holder.is_success)
      exceptions = exceptions_holder.exceptions
      self.assertEqual(len(exceptions), 1)
      exception_entry = exceptions[0]
      self.assertIsInstance(exception_entry.exception, CustomException)

# pytype: enable=wrong-arg-types

del BaseRunnerTestCase

if __name__ == "__main__":
  test_helper.run_pytest(__file__)