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

from typing import Iterable, Optional
import unittest
from unittest import mock

# //testing imports.
from unexpected_passes_common import builders
from unexpected_passes_common import constants
from unexpected_passes_common import data_types
from unexpected_passes_common import expectations
from unexpected_passes_common import queries
from unexpected_passes_common import unittest_utils as uu

# Protected access is allowed for unittests.
# pylint: disable=protected-access


class HelperMethodUnittest(unittest.TestCase):
  def testStripPrefixFromBuildIdValidId(self) -> None:
    self.assertEqual(queries._StripPrefixFromBuildId('build-1'), '1')

  def testStripPrefixFromBuildIdInvalidId(self) -> None:
    with self.assertRaises(AssertionError):
      queries._StripPrefixFromBuildId('build1')
    with self.assertRaises(AssertionError):
      queries._StripPrefixFromBuildId('build-1-2')

  def testConvertActualResultToExpectationFileFormatAbort(self) -> None:
    self.assertEqual(
      queries._ConvertActualResultToExpectationFileFormat('ABORT'), 'Timeout'
    )


class BigQueryQuerierInitUnittest(unittest.TestCase):
  def testInvalidNumSamples(self):
    """Tests that the number of samples is validated."""
    with self.assertRaises(AssertionError):
      uu.CreateGenericQuerier(num_samples=-1)

  def testDefaultSamples(self):
    """Tests that the number of samples is set to a default if not provided."""
    querier = uu.CreateGenericQuerier(num_samples=0)
    self.assertGreater(querier._num_samples, 0)


class GetBuilderGroupedQueryResultsUnittest(unittest.TestCase):
  def setUp(self):
    builders.ClearInstance()
    expectations.ClearInstance()
    uu.RegisterGenericBuildersImplementation()
    uu.RegisterGenericExpectationsImplementation()
    self._querier = uu.CreateGenericQuerier()

  def testUnknownBuilderType(self):
    """Tests behavior when an unknown builder type is provided."""
    with self.assertRaisesRegex(RuntimeError, 'Unknown builder type unknown'):
      for _ in self._querier.GetBuilderGroupedQueryResults('unknown', False):
        pass

  def testQueryRouting(self):
    """Tests that the correct query is used based on inputs."""
    with mock.patch.object(
      self._querier, '_GetPublicCiQuery', return_value='public_ci'
    ) as public_ci_mock:
      with mock.patch.object(
        self._querier, '_GetInternalCiQuery', return_value='internal_ci'
      ) as internal_ci_mock:
        with mock.patch.object(
          self._querier, '_GetPublicTryQuery', return_value='public_try'
        ) as public_try_mock:
          with mock.patch.object(
            self._querier, '_GetInternalTryQuery', return_value='internal_try'
          ) as internal_try_mock:
            all_mocks = [
              public_ci_mock,
              internal_ci_mock,
              public_try_mock,
              internal_try_mock,
            ]
            inputs = [
              (constants.BuilderTypes.CI, False, public_ci_mock),
              (constants.BuilderTypes.CI, True, internal_ci_mock),
              (constants.BuilderTypes.TRY, False, public_try_mock),
              (constants.BuilderTypes.TRY, True, internal_try_mock),
            ]
            for builder_type, internal_status, called_mock in inputs:
              for _ in self._querier.GetBuilderGroupedQueryResults(
                builder_type, internal_status
              ):
                pass
              for m in all_mocks:
                if m == called_mock:
                  m.assert_called_once()
                else:
                  m.assert_not_called()
              for m in all_mocks:
                m.reset_mock()

  def testNoResults(self):
    """Tests functionality if the query returns no results."""
    returned_builders = []
    with self.assertLogs(level='WARNING') as log_manager:
      with mock.patch.object(
        self._querier, '_GetPublicCiQuery', return_value=''
      ):
        for builder_name, _, _ in self._querier.GetBuilderGroupedQueryResults(
          constants.BuilderTypes.CI, False
        ):
          returned_builders.append(builder_name)
      for message in log_manager.output:
        if (
          'Did not get any results for builder type ci and internal status '
          'False. Depending on where tests are run and how frequently '
          'trybots are used for submission, this may be benign'
        ) in message:
          break
      else:
        self.fail('Did not find expected log message: %s' % log_manager.output)
      self.assertEqual(len(returned_builders), 0)

  def testHappyPath(self):
    """Tests functionality in the happy path."""
    self._querier.query_results = [
      uu.FakeQueryResult(
        builder_name='builder_a',
        id_='build-a',
        test_id='test_a',
        status='PASS',
        typ_tags=['linux', 'unknown_tag'],
        step_name='step_a',
      ),
      uu.FakeQueryResult(
        builder_name='builder_b',
        id_='build-b',
        test_id='test_b',
        status='FAIL',
        typ_tags=['win'],
        step_name='step_b',
      ),
    ]

    expected_results = [
      (
        'builder_a',
        [data_types.BaseResult('test_a', ('linux',), 'Pass', 'step_a', 'a')],
        None,
      ),
      (
        'builder_b',
        [data_types.BaseResult('test_b', ('win',), 'Failure', 'step_b', 'b')],
        None,
      ),
    ]

    results = []
    with mock.patch.object(self._querier, '_GetPublicCiQuery', return_value=''):
      for (
        builder_name,
        result_list,
        expectation_files,
      ) in self._querier.GetBuilderGroupedQueryResults(
        constants.BuilderTypes.CI, False
      ):
        results.append((builder_name, result_list, expectation_files))

    self.assertEqual(results, expected_results)

  def testHappyPathWithExpectationFiles(self):
    """Tests functionality in the happy path with expectation files provided."""
    self._querier.query_results = [
      uu.FakeQueryResult(
        builder_name='builder_a',
        id_='build-a',
        test_id='test_a',
        status='PASS',
        typ_tags=['linux', 'unknown_tag'],
        step_name='step_a',
      ),
      uu.FakeQueryResult(
        builder_name='builder_b',
        id_='build-b',
        test_id='test_b',
        status='FAIL',
        typ_tags=['win'],
        step_name='step_b',
      ),
    ]

    expected_results = [
      (
        'builder_a',
        [data_types.BaseResult('test_a', ('linux',), 'Pass', 'step_a', 'a')],
        list(set(['ef_a'])),
      ),
      (
        'builder_b',
        [data_types.BaseResult('test_b', ('win',), 'Failure', 'step_b', 'b')],
        list(set(['ef_b', 'ef_c'])),
      ),
    ]

    results = []
    with mock.patch.object(
      self._querier,
      '_GetRelevantExpectationFilesForQueryResult',
      side_effect=(['ef_a'], ['ef_b', 'ef_c']),
    ):
      with mock.patch.object(
        self._querier, '_GetPublicCiQuery', return_value=''
      ):
        for (
          builder_name,
          result_list,
          expectation_files,
        ) in self._querier.GetBuilderGroupedQueryResults(
          constants.BuilderTypes.CI, False
        ):
          results.append((builder_name, result_list, expectation_files))

    self.assertEqual(results, expected_results)


class FillExpectationMapForBuildersUnittest(unittest.TestCase):
  def setUp(self) -> None:
    self._querier = uu.CreateGenericQuerier()

    expectations.ClearInstance()
    uu.RegisterGenericExpectationsImplementation()

  def testErrorOnMixedBuilders(self) -> None:
    """Tests that providing builders of mixed type is an error."""
    builders_to_fill = [
      data_types.BuilderEntry('ci_builder', constants.BuilderTypes.CI, False),
      data_types.BuilderEntry('try_builder', constants.BuilderTypes.TRY, False),
    ]
    with self.assertRaises(AssertionError):
      self._querier.FillExpectationMapForBuilders(
        data_types.TestExpectationMap({}), builders_to_fill
      )

  def _runValidResultsTest(self, keep_unmatched_results: bool) -> None:
    self._querier = uu.CreateGenericQuerier(
      keep_unmatched_results=keep_unmatched_results
    )

    public_results = [
      uu.FakeQueryResult(
        builder_name='matched_builder',
        id_='build-build_id',
        test_id='foo',
        status='PASS',
        typ_tags=['win'],
        step_name='step_name',
      ),
      uu.FakeQueryResult(
        builder_name='unmatched_builder',
        id_='build-build_id',
        test_id='bar',
        status='PASS',
        typ_tags=[],
        step_name='step_name',
      ),
      uu.FakeQueryResult(
        builder_name='extra_builder',
        id_='build-build_id',
        test_id='foo',
        status='PASS',
        typ_tags=['win'],
        step_name='step_name',
      ),
    ]

    internal_results = [
      uu.FakeQueryResult(
        builder_name='matched_internal',
        id_='build-build_id',
        test_id='foo',
        status='PASS',
        typ_tags=['win'],
        step_name='step_name_internal',
      ),
      uu.FakeQueryResult(
        builder_name='unmatched_internal',
        id_='build-build_id',
        test_id='bar',
        status='PASS',
        typ_tags=[],
        step_name='step_name_internal',
      ),
    ]

    builders_to_fill = [
      data_types.BuilderEntry(
        'matched_builder', constants.BuilderTypes.CI, False
      ),
      data_types.BuilderEntry(
        'unmatched_builder', constants.BuilderTypes.CI, False
      ),
      data_types.BuilderEntry(
        'matched_internal', constants.BuilderTypes.CI, True
      ),
      data_types.BuilderEntry(
        'unmatched_internal', constants.BuilderTypes.CI, True
      ),
    ]

    expectation = data_types.Expectation(
      'foo', ['win'], 'RetryOnFailure', data_types.WildcardType.NON_WILDCARD
    )
    expectation_map = data_types.TestExpectationMap(
      {
        'foo': data_types.ExpectationBuilderMap(
          {
            expectation: data_types.BuilderStepMap(),
          }
        ),
      }
    )

    def PublicSideEffect():
      self._querier.query_results = public_results
      return ''

    def InternalSideEffect():
      self._querier.query_results = internal_results
      return ''

    with self.assertLogs(level='WARNING') as log_manager:
      with mock.patch.object(
        self._querier, '_GetPublicCiQuery', side_effect=PublicSideEffect
      ) as public_mock:
        with mock.patch.object(
          self._querier, '_GetInternalCiQuery', side_effect=InternalSideEffect
        ) as internal_mock:
          unmatched_results = self._querier.FillExpectationMapForBuilders(
            expectation_map, builders_to_fill
          )
          public_mock.assert_called_once()
          internal_mock.assert_called_once()

      for message in log_manager.output:
        if (
          'Did not find a matching builder for name extra_builder and '
          'internal status False. This is normal if the builder is no longer '
          'running tests (e.g. it was experimental).'
        ) in message:
          break
      else:
        self.fail('Did not find expected log message')

    stats = data_types.BuildStats()
    stats.AddPassedBuild(frozenset(['win']))
    expected_expectation_map = {
      'foo': {
        expectation: {
          'chromium/ci:matched_builder': {
            'step_name': stats,
          },
          'chrome/ci:matched_internal': {
            'step_name_internal': stats,
          },
        },
      },
    }
    self.assertEqual(expectation_map, expected_expectation_map)
    if keep_unmatched_results:
      self.assertEqual(
        unmatched_results,
        {
          'chromium/ci:unmatched_builder': [
            data_types.Result('bar', [], 'Pass', 'step_name', 'build_id'),
          ],
          'chrome/ci:unmatched_internal': [
            data_types.Result(
              'bar', [], 'Pass', 'step_name_internal', 'build_id'
            ),
          ],
        },
      )
    else:
      self.assertEqual(unmatched_results, {})

  def testValidResultsKeepUnmatched(self) -> None:
    """Tests behavior w/ valid results and keeping unmatched results."""
    self._runValidResultsTest(True)

  def testValidResultsDoNotKeepUnmatched(self) -> None:
    """Tests behavior w/ valid results and not keeping unmatched results."""
    self._runValidResultsTest(False)


class ProcessRowsForBuilderUnittest(unittest.TestCase):
  def setUp(self):
    self._querier = uu.CreateGenericQuerier()

  def testHappyPathWithExpectationFiles(self):
    """Tests functionality along the happy path with expectation files."""

    def SideEffect(row: queries.QueryResult) -> Optional[Iterable[str]]:
      if row.step_name == 'step_a1':
        return ['ef_a1']
      if row.step_name == 'step_a2':
        return ['ef_a2']
      if row.step_name == 'step_b':
        return ['ef_b1', 'ef_b2']
      raise RuntimeError('Unexpected row')

    rows = [
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-a',
        test_id='test_a',
        status='PASS',
        typ_tags=['linux', 'unknown_tag'],
        step_name='step_a1',
      ),
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-a',
        test_id='test_a',
        status='FAIL',
        typ_tags=['linux', 'unknown_tag'],
        step_name='step_a2',
      ),
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-b',
        test_id='test_b',
        status='FAIL',
        typ_tags=['win'],
        step_name='step_b',
      ),
    ]

    # Reversed order is expected since results are popped.
    expected_results = [
      data_types.BaseResult(
        test='test_b',
        tags=['win'],
        actual_result='Failure',
        step='step_b',
        build_id='b',
      ),
      data_types.BaseResult(
        test='test_a',
        tags=['linux'],
        actual_result='Failure',
        step='step_a2',
        build_id='a',
      ),
      data_types.BaseResult(
        test='test_a',
        tags=['linux'],
        actual_result='Pass',
        step='step_a1',
        build_id='a',
      ),
    ]

    with mock.patch.object(
      self._querier,
      '_GetRelevantExpectationFilesForQueryResult',
      side_effect=SideEffect,
    ):
      results, expectation_files = self._querier._ProcessRowsForBuilder(rows)
    self.assertEqual(results, expected_results)
    self.assertEqual(len(expectation_files), len(set(expectation_files)))
    self.assertEqual(
      set(expectation_files), set(['ef_a1', 'ef_a2', 'ef_b1', 'ef_b2'])
    )

  def testHappyPathNoneExpectation(self):
    """Tests functionality along the happy path with a None expectation file."""

    # A single None expectation file should cause the resulting return value to
    # become None.
    def SideEffect(row: queries.QueryResult) -> Optional[Iterable[str]]:
      if row.step_name == 'step_a1':
        return ['ef_a1']
      if row.step_name == 'step_a2':
        return ['ef_a2']
      return None

    rows = [
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-a',
        test_id='test_a',
        status='PASS',
        typ_tags=['linux', 'unknown_tag'],
        step_name='step_a1',
      ),
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-a',
        test_id='test_a',
        status='FAIL',
        typ_tags=['linux', 'unknown_tag'],
        step_name='step_a2',
      ),
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-b',
        test_id='test_b',
        status='FAIL',
        typ_tags=['win'],
        step_name='step_b',
      ),
    ]

    # Reversed order is expected since results are popped.
    expected_results = [
      data_types.BaseResult(
        test='test_b',
        tags=['win'],
        actual_result='Failure',
        step='step_b',
        build_id='b',
      ),
      data_types.BaseResult(
        test='test_a',
        tags=['linux'],
        actual_result='Failure',
        step='step_a2',
        build_id='a',
      ),
      data_types.BaseResult(
        test='test_a',
        tags=['linux'],
        actual_result='Pass',
        step='step_a1',
        build_id='a',
      ),
    ]

    with mock.patch.object(
      self._querier,
      '_GetRelevantExpectationFilesForQueryResult',
      side_effect=SideEffect,
    ):
      results, expectation_files = self._querier._ProcessRowsForBuilder(rows)
    self.assertEqual(results, expected_results)
    self.assertEqual(expectation_files, None)

  def testHappyPathSkippedResult(self):
    """Tests functionality along the happy path with a skipped result."""

    def SideEffect(row: queries.QueryResult) -> bool:
      if row.step_name == 'step_b':
        return True
      return False

    rows = [
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-a',
        test_id='test_a',
        status='PASS',
        typ_tags=['linux', 'unknown_tag'],
        step_name='step_a1',
      ),
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-a',
        test_id='test_a',
        status='FAIL',
        typ_tags=['linux', 'unknown_tag'],
        step_name='step_a2',
      ),
      uu.FakeQueryResult(
        builder_name='unused',
        id_='build-b',
        test_id='test_b',
        status='FAIL',
        typ_tags=['win'],
        step_name='step_b',
      ),
    ]

    # Reversed order is expected since results are popped.
    expected_results = [
      data_types.BaseResult(
        test='test_a',
        tags=['linux'],
        actual_result='Failure',
        step='step_a2',
        build_id='a',
      ),
      data_types.BaseResult(
        test='test_a',
        tags=['linux'],
        actual_result='Pass',
        step='step_a1',
        build_id='a',
      ),
    ]

    with mock.patch.object(
      self._querier, '_ShouldSkipOverResult', side_effect=SideEffect
    ):
      results, expectation_files = self._querier._ProcessRowsForBuilder(rows)
    self.assertEqual(results, expected_results)
    self.assertEqual(expectation_files, None)


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