# Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Unittests for build stages."""

from __future__ import print_function

import mock

from chromite.cbuildbot import cbuildbot_unittest
from chromite.cbuildbot.stages import artifact_stages
from chromite.cbuildbot.stages import generic_stages_unittest
from chromite.cbuildbot.stages import release_stages
from chromite.cbuildbot import failures_lib
from chromite.cbuildbot import results_lib
from chromite.lib import timeout_util

from chromite.cbuildbot.stages.generic_stages_unittest import patch

from chromite.lib.paygen import gspaths
from chromite.lib.paygen import paygen_build_lib


# pylint: disable=protected-access


class PaygenStageTest(generic_stages_unittest.AbstractStageTestCase,
                      cbuildbot_unittest.SimpleBuilderTestCase):
  """Test the PaygenStageStage."""

  BOT_ID = 'x86-mario-release'
  RELEASE_TAG = '0.0.1'

  SIGNER_RESULT = """
    { "status": { "status": "passed" }, "board": "link",
    "keyset": "link-mp-v4", "type": "recovery", "channel": "stable" }
    """

  INSNS_URLS_PER_CHANNEL = {
      'chan1': ['chan1_uri1', 'chan1_uri2'],
      'chan2': ['chan2_uri1'],
  }

  def setUp(self):
    self._Prepare()

  def ConstructStage(self):
    archive_stage = artifact_stages.ArchiveStage(self._run, self._current_board)
    return release_stages.PaygenStage(self._run, self._current_board,
                                      archive_stage)

  def testWaitForPushImageSuccess(self):
    """Test waiting for input from PushImage."""
    stage = self.ConstructStage()
    stage.board_runattrs.SetParallel(
        'instruction_urls_per_channel', self.INSNS_URLS_PER_CHANNEL)

    self.assertEqual(stage._WaitForPushImage(), self.INSNS_URLS_PER_CHANNEL)

  def testWaitForPushImageError(self):
    """Test WaitForPushImageError with an error output from pushimage."""
    stage = self.ConstructStage()
    stage.board_runattrs.SetParallel(
        'instruction_urls_per_channel', None)

    self.assertRaises(release_stages.MissingInstructionException,
                      stage._WaitForPushImage)

  def testWaitForSigningResultsSuccess(self):
    """Test that _WaitForSigningResults works when signing works."""
    results = ['chan1_uri1.json', 'chan1_uri2.json', 'chan2_uri1.json']

    with patch(release_stages.gs, 'GSContext') as mock_gs_ctx_init:
      mock_gs_ctx = mock_gs_ctx_init.return_value
      mock_gs_ctx.Cat.return_value = self.SIGNER_RESULT
      notifier = mock.Mock()

      stage = self.ConstructStage()
      stage._WaitForSigningResults(self.INSNS_URLS_PER_CHANNEL, notifier)

      self.assertEqual(notifier.mock_calls,
                       [mock.call('chan1'),
                        mock.call('chan2')])

      for result in results:
        mock_gs_ctx.Cat.assert_any_call(result)

  def testWaitForSigningResultsSuccessNothingSigned(self):
    """Test _WaitForSigningResults when there are no signed images."""
    with patch(release_stages.gs, 'GSContext') as mock_gs_ctx_init:
      mock_gs_ctx = mock_gs_ctx_init.return_value
      mock_gs_ctx.Cat.return_value = self.SIGNER_RESULT
      notifier = mock.Mock()

      stage = self.ConstructStage()
      stage._WaitForSigningResults({}, notifier)

      self.assertEqual(notifier.mock_calls, [])
      self.assertEqual(mock_gs_ctx.Cat.mock_calls, [])

  def testWaitForSigningResultsFailure(self):
    """Test _WaitForSigningResults when the signers report an error."""
    with patch(release_stages.gs, 'GSContext') as mock_gs_ctx_init:
      mock_gs_ctx = mock_gs_ctx_init.return_value
      mock_gs_ctx.Cat.return_value = """
          { "status": { "status": "failed" }, "board": "link",
            "keyset": "link-mp-v4", "type": "recovery", "channel": "stable" }
          """
      notifier = mock.Mock()

      stage = self.ConstructStage()

      self.assertRaisesStringifyable(
          release_stages.SignerFailure,
          stage._WaitForSigningResults,
          {'chan1': ['chan1_uri1']}, notifier)

      # Ensure we didn't notify anyone of success.
      self.assertEqual(notifier.mock_calls, [])
      self.assertEqual(mock_gs_ctx.Cat.mock_calls,
                       [mock.call('chan1_uri1.json')])

  def testWaitForSigningResultsTimeout(self):
    """Test that _WaitForSigningResults reports timeouts correctly."""
    with patch(release_stages.timeout_util, 'WaitForSuccess') as mock_wait:
      mock_wait.side_effect = timeout_util.TimeoutError
      notifier = mock.Mock()

      stage = self.ConstructStage()

      self.assertRaises(release_stages.SignerResultsTimeout,
                        stage._WaitForSigningResults,
                        {'chan1': ['chan1_uri1']}, notifier)

      self.assertEqual(notifier.mock_calls, [])

  def testCheckForResultsSuccess(self):
    """Test that _CheckForResults works when signing works."""
    with patch(release_stages.gs, 'GSContext') as mock_gs_ctx_init:
      mock_gs_ctx = mock_gs_ctx_init.return_value
      mock_gs_ctx.Cat.return_value = self.SIGNER_RESULT
      notifier = mock.Mock()

      stage = self.ConstructStage()
      self.assertTrue(
          stage._CheckForResults(mock_gs_ctx,
                                 self.INSNS_URLS_PER_CHANNEL,
                                 notifier))
      self.assertEqual(notifier.mock_calls,