#!/usr/bin/python

# Copyright (c) 2011 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."""

import mox
import os
import StringIO
import sys
import unittest

import constants
sys.path.append(constants.SOURCE_ROOT)
import chromite.buildbot.cbuildbot as cbuildbot
import chromite.buildbot.cbuildbot_config as config
import chromite.buildbot.cbuildbot_commands as commands
import chromite.buildbot.cbuildbot_stages as stages
import chromite.lib.cros_build_lib as cros_lib


class CBuildBotTest(mox.MoxTestBase):

  def setUp(self):
    mox.MoxTestBase.setUp(self)
    # Always stub RunCommmand out as we use it in every method.
    self.bot_id = 'x86-generic-pre-flight-queue'
    self.build_config = config.config[self.bot_id]

    self.options = self.mox.CreateMockAnything()
    self.options.buildroot = '.'
    self.options.resume = False

  def testNoResultCodeReturned(self):
    """Test a non-error run."""

    self.options.resume = True

    self.mox.StubOutWithMock(stages.BuilderStage.Results,
                             'RestoreCompletedStages')
    self.mox.StubOutWithMock(cbuildbot, 'RunBuildStages')
    self.mox.StubOutWithMock(stages.BuilderStage.Results,
                             'SaveCompletedStages')

    stages.BuilderStage.Results.RestoreCompletedStages(mox.IsA(file))

    cbuildbot.RunBuildStages(self.bot_id,
                             self.options,
                             self.build_config)

    stages.BuilderStage.Results.SaveCompletedStages(mox.IsA(file))

    self.mox.ReplayAll()

    cbuildbot.RunEverything(self.bot_id,
                            self.options,
                            self.build_config)

    self.mox.VerifyAll()

  # Verify bug 13035 is fixed.
  def testResultCodeReturned(self):
    """Verify that we return failure exit code on error."""

    self.options.resume = False

    self.mox.StubOutWithMock(stages.BuilderStage.Results,
                             'RestoreCompletedStages')
    self.mox.StubOutWithMock(cbuildbot, 'RunBuildStages')
    self.mox.StubOutWithMock(stages.BuilderStage.Results,
                             'SaveCompletedStages')

    cbuildbot.RunBuildStages(self.bot_id,
                             self.options,
                             self.build_config).AndRaise(
                                 Exception('Test Error'))

    stages.BuilderStage.Results.SaveCompletedStages(mox.IsA(file))

    self.mox.ReplayAll()

    self.assertRaises(
        SystemExit,
        lambda : cbuildbot.RunEverything(self.bot_id,
                                         self.options,
                                         self.build_config))

    self.mox.VerifyAll()


class BuilderStageTest(mox.MoxTestBase):

  def setUp(self):
    mox.MoxTestBase.setUp(self)
    # Always stub RunCommmand out as we use it in every method.
    self.bot_id = 'x86-generic-pre-flight-queue'
    self.build_config = config.config[self.bot_id]
    self.build_root = '/fake_root'
    self.url = 'fake_url'

    self.options = self.mox.CreateMockAnything()
    self.options.buildroot = self.build_root
    self.options.debug = False
    self.options.prebuilts = False
    self.options.tracking_branch = 'ooga_booga'
    self.options.clobber = False
    self.options.url = self.url
    self.options.buildnumber = 1234
    self.overlay = os.path.join(self.build_root,
                                'src/third_party/chromiumos-overlay')
    stages.BuilderStage.rev_overlays = [self.overlay]
    stages.BuilderStage.push_overlays = [self.overlay]
    self.mox.StubOutWithMock(os.path, 'isdir')

  def testGetPortageEnvVar(self):
    """Basic test case for _GetPortageEnvVar function."""
    self.mox.StubOutWithMock(cros_lib, 'OldRunCommand')
    envvar = 'EXAMPLE'
    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    cros_lib.OldRunCommand(mox.And(mox.IsA(list), mox.In(envvar)),
                           cwd='%s/src/scripts' % self.build_root,
                           redirect_stdout=True, enter_chroot=True,
                           error_ok=True).AndReturn('RESULT\n')
    self.mox.ReplayAll()
    stage = stages.BuilderStage(self.bot_id, self.options, self.build_config)
    result = stage._GetPortageEnvVar(envvar)
    self.mox.VerifyAll()
    self.assertEqual(result, 'RESULT')

  def testResolveOverlays(self):
    self.mox.StubOutWithMock(cros_lib, 'RunCommand')
    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    for _ in range(3):
      output_obj = cros_lib.CommandResult()
      output_obj.output = 'public1 public2\n'
      cros_lib.RunCommand(mox.And(mox.IsA(list), mox.In('--noprivate')),
                          redirect_stdout=True).AndReturn(output_obj)
      output_obj = cros_lib.CommandResult()
      output_obj.output = 'private1 private2\n'
      cros_lib.RunCommand(mox.And(mox.IsA(list), mox.In('--nopublic')), \
                          redirect_stdout=True).AndReturn(output_obj)
    self.mox.ReplayAll()
    stage = stages.BuilderStage(self.bot_id, self.options, self.build_config)
    public_overlays = ['public1', 'public2', self.overlay]
    private_overlays = ['private1', 'private2']
    self.assertEqual(stage._ResolveOverlays('public'), public_overlays)
    self.assertEqual(stage._ResolveOverlays('private'), private_overlays)
    self.assertEqual(stage._ResolveOverlays('both'),
                     public_overlays + private_overlays)
    self.mox.VerifyAll()


class SyncStageTest(BuilderStageTest):

  def setUp(self):
    mox.MoxTestBase.setUp(self)
    BuilderStageTest.setUp(self)
    self.mox.StubOutWithMock(commands, 'PreFlightRinse')

  def testFullSync(self):
    """Tests whether we can perform a full sync with a missing .repo folder."""
    self.mox.StubOutWithMock(commands, 'FullCheckout')

    os.path.isdir(self.build_root + '/.repo').AndReturn(False)
    os.path.isdir(self.build_root + '/.repo').AndReturn(False)
    commands.FullCheckout(self.build_root, self.options.tracking_branch,
                          url=self.url)
    os.path.isdir(self.overlay).AndReturn(True)

    self.mox.ReplayAll()
    stage = stages.SyncStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()

  def testIncrementalSync(self):
    """Tests whether we can perform a standard incremental sync."""
    self.mox.StubOutWithMock(commands, 'IncrementalCheckout')
    self.mox.StubOutWithMock(stages.BuilderStage, '_GetPortageEnvVar')

    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    commands.PreFlightRinse(self.build_root, self.build_config['board'],
                            self.options.tracking_branch, [self.overlay])
    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    stages.BuilderStage._GetPortageEnvVar(stages._FULL_BINHOST)
    commands.IncrementalCheckout(self.build_root)
    os.path.isdir(self.overlay).AndReturn(True)

    self.mox.ReplayAll()
    stage = stages.SyncStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()


class BuildBoardTest(BuilderStageTest):

  def setUp(self):
    mox.MoxTestBase.setUp(self)
    BuilderStageTest.setUp(self)

  def testFullBuild(self):
    """Tests whether correctly run make chroot and setup board for a full."""
    self.bot_id = 'x86-generic-full'
    self.build_config = config.config[self.bot_id]
    self.mox.StubOutWithMock(commands, 'MakeChroot')
    self.mox.StubOutWithMock(commands, 'SetupBoard')

    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    os.path.isdir(os.path.join(self.build_root, 'chroot')).AndReturn(True)
    commands.MakeChroot(self.build_root, self.build_config['chroot_replace'])
    os.path.isdir(os.path.join(self.build_root, 'chroot/build',
                               self.build_config['board'])).AndReturn(False)
    commands.SetupBoard(self.build_root, board=self.build_config['board'])

    self.mox.ReplayAll()
    stage = stages.BuildBoardStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()

  def testBinBuild(self):
    """Tests whether we skip un-necessary steps for a binary builder."""
    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    os.path.isdir(os.path.join(self.build_root, 'chroot')).AndReturn(True)
    os.path.isdir(os.path.join(self.build_root, 'chroot/build',
                               self.build_config['board'])).AndReturn(True)

    self.mox.ReplayAll()
    stage = stages.BuildBoardStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()

  def testBinBuildAfterClobber(self):
    """Tests whether we make chroot and board after a clobber."""
    self.mox.StubOutWithMock(commands, 'MakeChroot')
    self.mox.StubOutWithMock(commands, 'SetupBoard')

    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    os.path.isdir(os.path.join(self.build_root, 'chroot')).AndReturn(False)
    commands.MakeChroot(self.build_root, self.build_config['chroot_replace'])
    os.path.isdir(os.path.join(self.build_root, 'chroot/build',
                               self.build_config['board'])).AndReturn(False)
    commands.SetupBoard(self.build_root, board=self.build_config['board'])


    self.mox.ReplayAll()
    stage = stages.BuildBoardStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()


class BuildBoardTest(BuilderStageTest):

  def setUp(self):
    mox.MoxTestBase.setUp(self)
    BuilderStageTest.setUp(self)

  def testFullBuild(self):
    """Tests whether correctly run make chroot and setup board for a full."""
    self.bot_id = 'x86-generic-full'
    self.build_config = config.config[self.bot_id]
    self.mox.StubOutWithMock(commands, 'MakeChroot')
    self.mox.StubOutWithMock(commands, 'SetupBoard')

    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    os.path.isdir(os.path.join(self.build_root, 'chroot')).AndReturn(True)
    commands.MakeChroot(self.build_root, self.build_config['chroot_replace'])
    os.path.isdir(os.path.join(self.build_root, 'chroot/build',
                               self.build_config['board'])).AndReturn(False)
    commands.SetupBoard(self.build_root, board=self.build_config['board'])

    self.mox.ReplayAll()
    stage = stages.BuildBoardStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()

  def testBinBuild(self):
    """Tests whether we skip un-necessary steps for a binary builder."""
    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    os.path.isdir(os.path.join(self.build_root, 'chroot')).AndReturn(True)
    os.path.isdir(os.path.join(self.build_root, 'chroot/build',
                               self.build_config['board'])).AndReturn(True)

    self.mox.ReplayAll()
    stage = stages.BuildBoardStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()

  def testBinBuildAfterClobber(self):
    """Tests whether we make chroot and board after a clobber."""
    self.mox.StubOutWithMock(commands, 'MakeChroot')
    self.mox.StubOutWithMock(commands, 'SetupBoard')

    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    os.path.isdir(os.path.join(self.build_root, 'chroot')).AndReturn(False)
    commands.MakeChroot(self.build_root, self.build_config['chroot_replace'])
    os.path.isdir(os.path.join(self.build_root, 'chroot/build',
                               self.build_config['board'])).AndReturn(False)
    commands.SetupBoard(self.build_root, board=self.build_config['board'])


    self.mox.ReplayAll()
    stage = stages.BuildBoardStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()


class BuildTestsTest(BuilderStageTest):

  def setUp(self):
    mox.MoxTestBase.setUp(self)
    BuilderStageTest.setUp(self)
    self.fake_results_dir = '/tmp/fake_results_dir'

  def testFullTests(self):
    """Tests if full unit and cros_au_test_harness tests are run correctly."""
    self.bot_id = 'x86-generic-full'
    self.build_config = config.config[self.bot_id].copy()
    self.build_config['quick_unit'] = False
    self.build_config['quick_vm'] = False

    self.mox.StubOutWithMock(cros_lib, 'OldRunCommand')

    self.mox.StubOutWithMock(commands, 'RunUnitTests')
    self.mox.StubOutWithMock(commands, 'RunSmokeSuite')
    self.mox.StubOutWithMock(commands, 'RunAUTestSuite')
    self.mox.StubOutWithMock(commands, 'ArchiveTestResults')
    self.mox.StubOutWithMock(stages.TestStage, '_CreateTestRoot')

    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    stages.TestStage._CreateTestRoot().AndReturn(self.fake_results_dir)
    commands.RunUnitTests(self.build_root, full=True)
    commands.RunSmokeSuite(self.build_root, os.path.join(self.fake_results_dir,
                                                         'smoke_results'))
    commands.RunAUTestSuite(self.build_root,
                            self.build_config['board'],
                            os.path.join(self.fake_results_dir,
                                         'au_test_harness'),
                            full=True)
    commands.ArchiveTestResults(self.build_root, self.fake_results_dir)

    self.mox.ReplayAll()
    stage = stages.TestStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()

  def testQuickTests(self):
    """Tests if quick unit and cros_au_test_harness tests are run correctly."""
    self.bot_id = 'x86-generic-full'
    self.build_config = config.config[self.bot_id].copy()
    self.build_config['quick_unit'] = True
    self.build_config['quick_vm'] = True

    self.mox.StubOutWithMock(commands, 'RunUnitTests')
    self.mox.StubOutWithMock(commands, 'RunSmokeSuite')
    self.mox.StubOutWithMock(commands, 'RunAUTestSuite')
    self.mox.StubOutWithMock(commands, 'ArchiveTestResults')
    self.mox.StubOutWithMock(stages.TestStage, '_CreateTestRoot')

    os.path.isdir(self.build_root + '/.repo').AndReturn(True)
    stages.TestStage._CreateTestRoot().AndReturn(self.fake_results_dir)
    commands.RunUnitTests(self.build_root, full=False)
    commands.RunSmokeSuite(self.build_root, os.path.join(self.fake_results_dir,
                                                         'smoke_results'))
    commands.RunAUTestSuite(self.build_root,
                            self.build_config['board'],
                            os.path.join(self.fake_results_dir,
                                         'au_test_harness'),
                            full=False)
    commands.ArchiveTestResults(self.build_root, self.fake_results_dir)

    self.mox.ReplayAll()
    stage = stages.TestStage(self.bot_id, self.options, self.build_config)
    stage.Run()
    self.mox.VerifyAll()


class BuildStagesResultsTest(unittest.TestCase):

  def setUp(self):
    unittest.TestCase.setUp(self)
    # Always stub RunCommmand out as we use it in every method.
    self.bot_id = 'x86-generic-pre-flight-queue'
    self.build_config = config.config[self.bot_id]
    self.build_root = '/fake_root'
    self.url = 'fake_url'

    # Create a class to hold 
    class Options:
      pass

    self.options = Options()
    self.options.buildroot = self.build_root
    self.options.debug = False
    self.options.prebuilts = False
    self.options.tracking_branch = 'ooga_booga'
    self.options.clobber = False
    self.options.url = self.url
    self.options.buildnumber = 1234

    self.failException = Exception("FailStage needs to fail.")


  def _runStages(self):
    """Run a couple of stages so we can capture the results"""
    class PassStage(stages.BuilderStage):
      """PassStage always works"""
      pass

    class FailStage(stages.BuilderStage):
      """FailStage always throws an exception"""

      def _PerformStage(localSelf):
        """Throw the exception to make us fail."""
        raise self.failException

    # Run two stages
    PassStage(self.bot_id, self.options, self.build_config).Run()

    self.assertRaises(
        Exception,
        FailStage(self.bot_id, self.options, self.build_config).Run)

  def testRunStages(self):
    """Run some stages and verify the captured results"""

    stages.BuilderStage.Results.Clear()
    self.assertEqual(stages.BuilderStage.Results.Get(), [])

    self._runStages()

    # Verify that the results are what we expect.
    expectedResults = [
        ('Pass', stages.BuilderStage.Results.SUCCESS),
        ('Fail', self.failException)]

    actualResults = stages.BuilderStage.Results.Get()
    
    # Break out the asserts to be per item to make debugging easier
    self.assertEqual(len(expectedResults), len(actualResults))
    for i in xrange(len(expectedResults)):
      self.assertEqual(expectedResults[i], actualResults[i])

  def testStagesReport(self):
    """Tests Stage reporting."""

    # Store off a known set of results and generate a report
    stages.BuilderStage.Results.Clear()
    stages.BuilderStage.Results.Record('Pass',
                                       stages.BuilderStage.Results.SKIPPED)
    stages.BuilderStage.Results.Record('Pass2',
                                       stages.BuilderStage.Results.SUCCESS)
    stages.BuilderStage.Results.Record('Fail', self.failException)
    stages.BuilderStage.Results.Record(
        'FailRunCommand',
        cros_lib.RunCommandError(
            'Command "/bin/false /nosuchdir" failed.\n',
            ['/bin/false', '/nosuchdir']))
    stages.BuilderStage.Results.Record(
        'FailOldRunCommand',
        cros_lib.RunCommandException(
            'Command "[\'/bin/false\', \'/nosuchdir\']" failed.\n',
            ['/bin/false', '/nosuchdir']))

    results = StringIO.StringIO()

    stages.BuilderStage.Results.Report(results)

    expectedResults = (
        "************************************************************\n"
        "** Stage Results\n"
        "************************************************************\n"
        "** Pass previously completed\n"
        "************************************************************\n"
        "** Pass2\n"
        "************************************************************\n"
        "** Fail failed with Exception\n"
        "************************************************************\n"
        "** FailRunCommand failed in /bin/false\n"
        "************************************************************\n"
        "** FailOldRunCommand failed in /bin/false\n"
        "************************************************************\n")

    expectedLines = expectedResults.split('\n')
    actualLines = results.getvalue().split('\n')

    # Break out the asserts to be per item to make debugging easier
    self.assertEqual(len(expectedLines), len(actualLines))
    for i in xrange(len(expectedLines)):
      self.assertEqual(expectedLines[i], actualLines[i])

  def testSaveCompletedStages(self):
    """Tests that we can save out completed stages."""

    # Run this again to make sure we have the expected results stored
    stages.BuilderStage.Results.Clear()
    stages.BuilderStage.Results.Record('Pass',
                                       stages.BuilderStage.Results.SUCCESS)

    stages.BuilderStage.Results.Record('Fail',
                                       self.failException)

    stages.BuilderStage.Results.Record('Pass2',
                                       stages.BuilderStage.Results.SUCCESS)

    saveFile = StringIO.StringIO()
    stages.BuilderStage.Results.SaveCompletedStages(saveFile)
    self.assertEqual(saveFile.getvalue(), 'Pass\n')

  def testRestoreCompletedStages(self):
    """Tests that we can read in completed stages."""

    stages.BuilderStage.Results.Clear()
    stages.BuilderStage.Results.RestoreCompletedStages(
        StringIO.StringIO('Pass\n'))

    self.assertEqual(stages.BuilderStage.Results.GetPrevious(), ['Pass'])

  def testRunAfterRestore(self):
    """Tests that we skip previously completed stages."""

    # Fake stages.BuilderStage.Results.RestoreCompletedStages
    stages.BuilderStage.Results.Clear()
    stages.BuilderStage.Results.RestoreCompletedStages(
        StringIO.StringIO('Pass\n'))

    self._runStages()

    # Verify that the results are what we expect.
    expectedResults = [
        ('Pass', stages.BuilderStage.Results.SKIPPED),
        ('Fail', self.failException)]

    actualResults = stages.BuilderStage.Results.Get()

    # Break out the asserts to be per item to make debugging easier
    self.assertEqual(len(expectedResults), len(actualResults))
    for i in xrange(len(expectedResults)):
      self.assertEqual(expectedResults[i], actualResults[i])

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