# Copyright 2018 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 the security_test_image script."""

import os

from chromite.lib import cros_build_lib
from chromite.lib import cros_test_lib
from chromite.lib import image_lib
from chromite.scripts import security_test_image


class SecurityTestImageTest(cros_test_lib.MockTempDirTestCase):
    """Security test image script tests."""

    def setUp(self) -> None:
        D = cros_test_lib.Directory
        filesystem = (
            D("board", ("recovery_image.bin",)),
            "other_image.bin",
        )
        cros_test_lib.CreateOnDiskHierarchy(self.tempdir, filesystem)

    def testParseArgs(self) -> None:
        """Argument parsing tests."""
        # pylint: disable=protected-access
        # Test no arguments.
        with self.assertRaises(SystemExit):
            security_test_image._ParseArgs([])

        # Test board is set but not used when we have the full image path.
        self.PatchObject(
            cros_build_lib, "GetDefaultBoard", return_value="board"
        )
        opts = security_test_image._ParseArgs(
            ["--image", os.path.join(self.tempdir, "other_image.bin")]
        )
        self.assertEqual("board", opts.board)
        self.assertEqual(
            opts.image, os.path.join(self.tempdir, "other_image.bin")
        )

        # Test the board is fetched and used when using the default image
        # basename.
        self.PatchObject(
            image_lib,
            "GetLatestImageLink",
            return_value=os.path.join(self.tempdir, "board"),
        )
        opts = security_test_image._ParseArgs([])
        self.assertEqual("board", opts.board)
        self.assertEqual(
            opts.image, os.path.join(self.tempdir, "board/recovery_image.bin")
        )