# Copyright 2012 The ChromiumOS Authors
# 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 unittest import mock

from chromite.cbuildbot import cbuildbot_unittest
from chromite.cbuildbot.stages import generic_stages_unittest
from chromite.cbuildbot.stages import release_stages
from chromite.cbuildbot.stages.generic_stages_unittest import patch
from chromite.lib import config_lib
from chromite.lib import constants
from chromite.lib import failures_lib
from chromite.lib import parallel
from chromite.lib import parallel_unittest
from chromite.lib import results_lib
from chromite.lib import timeout_util
from chromite.lib.buildstore import FakeBuildStore
from chromite.lib.paygen import gspaths
from chromite.lib.paygen import paygen_build_lib


# pylint: disable=protected-access


class SigningStageTest(
    generic_stages_unittest.AbstractStageTestCase,
    cbuildbot_unittest.SimpleBuilderTestCase,
):
    """Test the SigningStage."""

    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()
        self.buildstore = FakeBuildStore()

    def ConstructStage(self):
        return release_stages.SigningStage(
            self._run, self.buildstore, self._current_board
        )

    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.WaitUntilReady(), True)
        self.assertEqual(
            stage.instruction_urls_per_channel, 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.assertEqual(stage.WaitUntilReady(), False)

    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 testWaitForSigningResultsSuccessNoNotifier(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

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

            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, [mock.call("chan1"), mock.call("chan2")]
            )

    def testCheckForResultsSuccessNoChannels(self):
        """Test that _CheckForResults works when there is nothing to check for."""
        with patch(release_stages.gs, "GSContext") as mock_gs_ctx_init:
            mock_gs_ctx = mock_gs_ctx_init.return_value
            notifier = mock.Mock()

            stage = self.ConstructStage()

            # Ensure we find that we are ready if there are no channels to look for.
            self.assertTrue(stage._CheckForResults(mock_gs_ctx, {}, notifier))

            # Ensure we didn't contact GS while checking for no channels.
            self.assertFalse(mock_gs_ctx.Cat.called)
            self.assertEqual(notifier.mock_calls, [])

    def testCheckForResultsPartialComplete(self):
        """Verify _CheckForResults handles partial signing results."""

        def catChan2Success(url):
            if url.startswith("chan2"):
                return self.SIGNER_RESULT
            else:
                raise release_stages.gs.GSNoSuchKey()

        with patch(release_stages.gs, "GSContext") as mock_gs_ctx_init:
            mock_gs_ctx = mock_gs_ctx_init.return_value
            mock_gs_ctx.Cat.side_effect = catChan2Success
            notifier = mock.Mock()

            stage = self.ConstructStage()
            self.assertFalse(
                stage._CheckForResults(
                    mock_gs_ctx, self.INSNS_URLS_PER_CHANNEL, notifier
                )
            )
            self.assertEqual(
                stage.signing_results,
                {
                    "chan1": {},
                    "chan2": {
                        "chan2_uri1.json": {
                            "board": "link",
                            "channel": "stable",
                            "keyset": "link-mp-v4",
                            "status": {"status": "passed"},
                            "type": "recovery",
                        }
                    },
                },
            )
            self.assertEqual(notifier.mock_calls, [mock.call("chan2")])

    def testCheckForResultsUnexpectedJson(self):
        """Verify _CheckForResults handles unexpected Json values."""
        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 = "{}"
            notifier = mock.Mock()

            stage = self.ConstructStage()
            self.assertFalse(
                stage._CheckForResults(
                    mock_gs_ctx, self.INSNS_URLS_PER_CHANNEL, notifier
                )
            )
            self.assertEqual(stage.signing_results, {"chan1": {}, "chan2": {}})
            self.assertEqual(notifier.mock_calls, [])

    def testCheckForResultsMalformedJson(self):
        """Verify _CheckForResults handles unexpected Json values."""
        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 = "{"
            notifier = mock.Mock()

            stage = self.ConstructStage()
            self.assertFalse(
                stage._CheckForResults(
                    mock_gs_ctx, self.INSNS_URLS_PER_CHANNEL, notifier
                )
            )
            self.assertEqual(stage.signing_results, {"chan1": {}, "chan2": {}})
            self.assertEqual(notifier.mock_calls, [])

    def testCheckForResultsNoResult(self):
        """Verify _CheckForResults handles missing signer results."""
        with patch(release_stages.gs, "GSContext") as mock_gs_ctx_init:
            mock_gs_ctx = mock_gs_ctx_init.return_value
            mock_gs_ctx.Cat.side_effect = release_stages.gs.GSNoSuchKey
            notifier = mock.Mock()

            stage = self.ConstructStage()
            self.assertFalse(
                stage._CheckForResults(
                    mock_gs_ctx, self.INSNS_URLS_PER_CHANNEL, notifier
                )
            )
            self.assertEqual(stage.signing_results, {"chan1": {}, "chan2": {}})
            self.assertEqual(notifier.mock_calls, [])

    def testCheckForResultsFailed(self):
        """Verify _CheckForResults handles missing signer results."""
        with patch(release_stages.gs, "GSContext") as mock_gs_ctx_init:
            mock_gs_ctx = mock_gs_ctx_init.return_value
            mock_gs_ctx.Cat.side_effect = release_stages.gs.GSNoSuchKey
            notifier = mock.Mock()

            stage = self.ConstructStage()
            self.assertFalse(
                stage._CheckForResults(
                    mock_gs_ctx, self.INSNS_URLS_PER_CHANNEL, notifier
                )
            )
            self.assertEqual(stage.signing_results, {"chan1": {}, "chan2": {}})
            self.assertEqual(notifier.mock_calls, [])

    def testPerformStageSuccess(self):
        """Test that SigningStage works when signing works."""
        stage = self.ConstructStage()
        stage.instruction_urls_per_channel = self.INSNS_URLS_PER_CHANNEL
        self.PatchObject(stage, "_WaitForSigningResults")

        stage.PerformStage()

        # Verify that we send the right notifications.
        result = stage.board_runattrs.GetParallel(
            "signed_images_ready", timeout=0
        )
        self.assertEqual(result, ["chan1", "chan2"])


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

    # We use a variant board to make sure the '_' is translated to '-'.
    BOT_ID = "beaglebone_servo-release"
    RELEASE_TAG = "0.0.1"

    def setUp(self):
        self._Prepare()
        self.buildstore = FakeBuildStore()

        # This method fetches a file from GS, mock it out.
        self.validateMock = self.PatchObject(
            paygen_build_lib, "ValidateBoardConfig"
        )

        instanceMock = mock.MagicMock()
        self.paygenBuildMock = self.PatchObject(
            paygen_build_lib, "PaygenBuild", return_value=instanceMock
        )

        self.payload_config1 = mock.MagicMock()
        self.payload_config1.payload_type = paygen_build_lib.PAYLOAD_TYPE_FSI
        self.payload_config1.applicable_models = ["model1", "model3"]
        self.payload_config2 = mock.MagicMock()
        self.payload_config2.payload_type = paygen_build_lib.PAYLOAD_TYPE_OMAHA
        self.payload_config2.applicable_models = ["model2", "model3"]
        self.payload_config3 = mock.MagicMock()
        self.payload_config3.payload_type = paygen_build_lib.PAYLOAD_TYPE_OMAHA
        self.payload_config4 = mock.MagicMock()
        self.payload_config4.payload_type = paygen_build_lib.PAYLOAD_TYPE_N2N
        self.payload_config5 = mock.MagicMock()
        self.payload_config5.payload_type = (
            paygen_build_lib.PAYLOAD_TYPE_STEPPING_STONE
        )

        instanceMock.CreatePayloads.side_effect = iter(
            [
                (
                    "foo-suite-name",
                    "foo-archive-board",
                    "foo-archive-build",
                    [
                        self.payload_config1,
                        self.payload_config2,
                        self.payload_config3,
                        self.payload_config4,
                        self.payload_config5,
                    ],
                )
            ]
        )

    # pylint: disable=arguments-differ
    def ConstructStage(self, channels=None):
        return release_stages.PaygenStage(
            self._run, self.buildstore, self._current_board, channels=channels
        )

    def testWaitUntilReadSigning(self):
        """Test that PaygenStage works when signing works."""
        stage = self.ConstructStage()
        stage.board_runattrs.SetParallel(
            "signed_images_ready", ["stable", "beta"]
        )

        self.assertEqual(stage.WaitUntilReady(), True)
        self.assertEqual(stage.channels, ["stable", "beta"])

    def testWaitUntilReadSigningFailure(self):
        """Test that PaygenStage works when signing works."""
        stage = self.ConstructStage()
        stage.board_runattrs.SetParallel("signed_images_ready", None)

        self.assertEqual(stage.WaitUntilReady(), False)

    def testWaitUntilReadSigningEmpty(self):
        """Test that PaygenStage works when signing works."""
        stage = self.ConstructStage()
        stage.board_runattrs.SetParallel("signed_images_ready", [])

        self.assertEqual(stage.WaitUntilReady(), True)
        self.assertEqual(stage.channels, [])

    def testPerformStageSuccess(self):
        """Test that PaygenStage works when signing works."""
        with patch(release_stages.PaygenStage, "_RunPaygenInProcess") as runner:
            channels = ["stable", "beta"]
            stage = self.ConstructStage(channels=channels)

            with parallel_unittest.ParallelMock():
                stage.PerformStage()

            metadata_channels = self._run.attrs.metadata.GetValue("channels")
            self.assertEqual(",".join(channels), metadata_channels)

            # Verify that we validate with the board name in release name space.
            self.assertEqual(
                self.validateMock.call_args_list,
                [mock.call("beaglebone-servo")],
            )

            # Verify that we queue up work.
            runner.assert_has_calls(
                [
                    mock.call(
                        stage,
                        "stable",
                        "beaglebone-servo",
                        "0.0.1",
                        False,
                        False,
                        False,
                    ),
                    mock.call(
                        stage,
                        "beta",
                        "beaglebone-servo",
                        "0.0.1",
                        False,
                        False,
                        False,
                    ),
                ]
            )

    def testPerformStageNoChannels(self):
        """Test that PaygenStage works when signing works."""
        with patch(release_stages.PaygenStage, "_RunPaygenInProcess") as runner:
            stage = self.ConstructStage(channels=[])

            with parallel_unittest.ParallelMock():
                stage.PerformStage()

            # Verify that we did not queue up work.
            runner.assert_not_called()

    def testPerformStageTrybot(self):
        """Test the PerformStage alternate behavior for trybot runs."""
        with patch(release_stages.PaygenStage, "_RunPaygenInProcess") as runner:
            # The stage is constructed differently for trybots, so don't use
            # ConstructStage.
            stage = self.ConstructStage(channels=["foo", "bar"])
            with parallel_unittest.ParallelMock():
                stage.PerformStage()

            # Notice that we didn't put anything in _wait_for_channel_signing, but
            # still got results right away.
            runner.assert_has_calls(
                [
                    mock.call(
                        stage,
                        "foo",
                        "beaglebone-servo",
                        "0.0.1",
                        False,
                        False,
                        False,
                    ),
                    mock.call(
                        stage,
                        "bar",
                        "beaglebone-servo",
                        "0.0.1",
                        False,
                        False,
                        False,
                    ),
                ]
            )

    def testPerformStageUnknownBoard(self):
        """Test that PaygenStage exits when an unknown board is specified."""
        self._current_board = "unknown-board-name"

        # Setup a board validation failure.
        badBoardException = paygen_build_lib.BoardNotConfigured(
            self._current_board
        )
        self.validateMock.side_effect = badBoardException

        stage = self.ConstructStage()

        with self.assertRaises(release_stages.PaygenNoPaygenConfigForBoard):
            stage.PerformStage()

    def testRunPaygenInProcess(self):
        """Test that _RunPaygenInProcess works in the simple case."""
        # Have to patch and verify that the PaygenTestStage is created.
        stage = self.ConstructStage()

        with patch(paygen_build_lib, "ScheduleAutotestTests") as sched_tests:
            # Call the method under test.
            stage._RunPaygenInProcess(
                "foo", "foo-board", "foo-version", True, False, False
            )
            # Ensure that PaygenTestStage is created and schedules the test suite
            # with the correct arguments.
            sched_tests.assert_called_once_with(
                "foo-suite-name",
                "foo-archive-board",
                None,
                "foo-archive-build",
                [
                    self.payload_config1,
                    self.payload_config2,
                    self.payload_config3,
                    self.payload_config4,
                    self.payload_config5,
                ],
            )

        # Ensure arguments are properly converted and passed along.
        self.paygenBuildMock.assert_called_with(
            gspaths.Build(
                version="foo-version",
                board="foo-board",
                channel="foo-channel",
                bucket=gspaths.ChromeosReleases.BUCKET,
            ),
            mock.ANY,
            work_dir=mock.ANY,
            site_config=stage._run.site_config,
            dry_run=True,
            skip_delta_payloads=False,
        )

    def testRunPaygenInProcessInSkylab(self):
        """Test that _RunPaygenInProcess works in Skylab."""
        self._run.config.enable_skylab_hw_tests = True
        # Have to patch and verify that the PaygenTestStage is created.
        stage = self.ConstructStage()

        with patch(paygen_build_lib, "ScheduleAutotestTests") as sched_tests:
            # Call the method under test.
            stage._RunPaygenInProcess(
                "foo", "foo-board", "foo-version", True, False, False
            )
            # Ensure that PaygenTestStage is created and schedules the test suite
            # with the correct arguments.
            sched_tests.assert_called_once_with(
                "foo-suite-name",
                "foo-archive-board",
                None,
                "foo-archive-build",
                [
                    self.payload_config1,
                    self.payload_config2,
                    self.payload_config3,
                    self.payload_config4,
                    self.payload_config5,
                ],
            )

    def testRunPaygenInProcessComplex(self):
        """Test that _RunPaygenInProcess with arguments that are more unusual."""
        # Call the method under test.
        # Use release tools channel naming, and a board name including a variant.
        stage = self.ConstructStage()
        stage._RunPaygenInProcess(
            "foo-channel", "foo-board-variant", "foo-version", True, True, True
        )

        # Ensure arguments are properly converted and passed along.
        self.paygenBuildMock.assert_called_with(
            gspaths.Build(
                version="foo-version",
                board="foo-board-variant",
                channel="foo-channel",
                bucket=gspaths.ChromeosReleases.BUCKET,
            ),
            mock.ANY,
            dry_run=True,
            work_dir=mock.ANY,
            site_config=stage._run.site_config,
            skip_delta_payloads=True,
        )

    def testRunPaygenInProcessWithUnifiedBuild(self):
        self._run.config.models = [
            config_lib.ModelTestConfig("model1", "model1"),
            config_lib.ModelTestConfig("model2", "board", ["au"]),
        ]

        # Have to patch and verify that the PaygenTestStage is created.
        stage = self.ConstructStage()

        with patch(paygen_build_lib, "ScheduleAutotestTests") as sched_tests:
            # Call the method under test.
            stage._RunPaygenInProcess(
                "foo", "foo-board", "foo-version", True, False, False
            )
            # Ensure that the first model from the unified build was selected
            # as the platform to be tested
            sched_tests.assert_called_once_with(
                "foo-suite-name",
                "board",
                "model2",
                "foo-archive-build",
                [
                    self.payload_config1,
                    self.payload_config2,
                    self.payload_config3,
                    self.payload_config4,
                    self.payload_config5,
                ],
            )

    def testRunPaygenInProcessWithUnifiedBuildInSkylab(self):
        """Test that _RunPaygenInProcess works for unibuild in Skylab."""
        self._run.config.enable_skylab_hw_tests = True
        self._run.config.models = [
            config_lib.ModelTestConfig(
                "model2", "board", ["au"], enable_skylab=True
            ),
        ]

        # Have to patch and verify that the PaygenTestStage is created.
        stage = self.ConstructStage()

        with patch(paygen_build_lib, "ScheduleAutotestTests") as sched_tests:
            # Call the method under test.
            stage._RunPaygenInProcess(
                "foo", "foo-board", "foo-version", True, False, False
            )
            # Ensure that the first model from the unified build was selected
            # as the platform to be tested
            sched_tests.assert_called_once_with(
                "foo-suite-name",
                "board",
                "model2",
                "foo-archive-build",
                [
                    self.payload_config1,
                    self.payload_config2,
                    self.payload_config3,
                    self.payload_config4,
                    self.payload_config5,
                ],
            )

    def testRunPaygenInParallelWithUnifiedBuild(self):
        # payload_config1 defines applicable_models as model1 and model3.
        # model3 does not have au enabled but gets scheduled since it has type FSI.
        # payload_config2 defines applicable_models as model2 and model3.
        # model3 does not get scheduled since config2 has type OMAHA.
        # payload_config3 has type OMAHA with no applicable models so doesn't get
        # scheduled.
        # payload_config4 has type N2N and which never have applicable_models but
        # should get scheduled on all ['au'] models.
        # payload_config5 has type STEPPING_STONE which also never has
        # applicable models but should get scheduled on all ['au'] models too.
        self._run.config.models = [
            config_lib.ModelTestConfig("model1", "model1", ["au"]),
            config_lib.ModelTestConfig("model2", "model1", ["au"]),
            config_lib.ModelTestConfig("model3", "model1"),
        ]

        # Have to patch and verify that the PaygenTestStage is created.
        stage = self.ConstructStage()

        with patch(parallel, "RunParallelSteps") as parallel_tests:
            stage._RunPaygenInProcess(
                "foo", "foo-board", "foo-version", True, False, False
            )
            # 2 tests scheduled for FSI, 1 test scheduled for OMAHA,
            # 2 tests scheduled for N2N, 2 tests scheduled for STEPPING_STONE.
            parallel_tests.assert_called_once_with([mock.ANY] * 7)

    def testPayloadBuildSetCorrectly(self):
        """Test that payload build is passed correctly to PaygenBuild."""
        stage = self.ConstructStage()
        self.PatchObject(paygen_build_lib, "ScheduleAutotestTests")

        # Call the method under test.
        stage._RunPaygenInProcess(
            "foo", "foo-board", "foo-version", False, False, False
        )

        # Ensure arguments are properly converted and passed along.
        build = gspaths.Build(
            version="foo-version",
            board="foo-board",
            channel="foo-channel",
            bucket=gspaths.ChromeosReleases.BUCKET,
        )
        self.paygenBuildMock.assert_called_with(
            build,
            build,
            work_dir=mock.ANY,
            site_config=stage._run.site_config,
            dry_run=False,
            skip_delta_payloads=False,
        )

    def testTestPayloadBuildSetCorrectly(self):
        """Test that test payload build is passed correctly to PaygenBuild."""
        stage = self.ConstructStage()
        self.PatchObject(paygen_build_lib, "ScheduleAutotestTests")

        # Call the method under test.
        stage._RunPaygenInProcess(
            "foo", "foo-board", "foo-version", True, False, False
        )

        # Ensure arguments are properly converted and passed along.
        build = gspaths.Build(
            version="foo-version",
            board="foo-board",
            channel="foo-channel",
            bucket=gspaths.ChromeosReleases.BUCKET,
        )
        test_payload_build = gspaths.Build(build)
        test_payload_build.bucket = gspaths.ChromeosReleases.TEST_BUCKET

        self.paygenBuildMock.assert_called_with(
            build,
            test_payload_build,
            work_dir=mock.ANY,
            site_config=stage._run.site_config,
            dry_run=True,
            skip_delta_payloads=False,
        )


class PaygenBuildStageTest(
    generic_stages_unittest.AbstractStageTestCase,
    cbuildbot_unittest.SimpleBuilderTestCase,
):
    """Test the PaygenBuild stage."""

    # We use a variant board to make sure the '_' is translated to '-'.
    BOT_ID = "beaglebone_servo-release"
    RELEASE_TAG = "0.0.1"

    def setUp(self):
        self._Prepare()
        self.buildstore = FakeBuildStore()

    # pylint: disable=arguments-differ
    def ConstructStage(self):
        return release_stages.PaygenBuildStage(
            self._run,
            self.buildstore,
            board=self._current_board,
            channel="foochan",
            version="foo-version",
            debug=True,
            skip_testing=False,
            skip_delta_payloads=False,
        )

    def testStageName(self):
        """See if the stage name is correctly formed."""
        stage = self.ConstructStage()
        self.assertEqual(stage.name, "PaygenBuildFoochan")


class PaygenTestStageTest(
    generic_stages_unittest.AbstractStageTestCase,
    cbuildbot_unittest.SimpleBuilderTestCase,
):
    """Test the PaygenTestStage stage."""

    # We use a variant board to make sure the '_' is translated to '-'.
    BOT_ID = "beaglebone_servo-release"
    RELEASE_TAG = "0.0.1"

    def setUp(self):
        self._Prepare()
        self.buildstore = FakeBuildStore()

    # pylint: disable=arguments-differ
    def ConstructStage(self):
        return release_stages.PaygenTestStage(
            builder_run=self._run,
            buildstore=self.buildstore,
            suite_name="foo-test-suite",
            board=self._current_board,
            model=self._current_board,
            lab_board_name=self._current_board,
            # The PaygenBuild stage will add the '-channel' suffix to the channel
            # when converting to release tools naming.
            channel="foochan-channel",
            build="foo-version",
            debug=True,
            payload_test_configs=[],
            test_env=constants.ENV_AUTOTEST,
        )

    def testStageName(self):
        """See if the stage name is correctly formed."""
        stage = self.ConstructStage()
        self.assertEqual(stage.name, "PaygenTestFoochan [beaglebone_servo]")

    def testPerformStageTestLabFail(self):
        """Test that exception from RunHWTestSuite are properly handled."""
        with patch(paygen_build_lib, "ScheduleAutotestTests") as sched_tests:
            sched_tests.side_effect = failures_lib.TestLabFailure

            stage = self.ConstructStage()

            with patch(stage, "_HandleExceptionAsWarning") as warning_handler:
                warning_handler.return_value = (
                    results_lib.Results.FORGIVEN,
                    "description",
                    0,
                )

                stage.Run()

                # This proves the exception was turned into a warning.
                self.assertTrue(warning_handler.called)