xref: /aosp_15_r20/external/toolchain-utils/crosperf/download_images_unittest.py (revision 760c253c1ed00ce9abd48f8546f08516e57485fe)

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# Copyright 2019 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""Download image unittest."""


import os
import re
import unittest
import unittest.mock as mock

from cros_utils import command_executer
from cros_utils import logger
import download_images
import test_flag


MOCK_LOGGER = logger.GetLogger(log_dir="", mock=True)


class RegexMatcher:
    """A regex matcher, for passing to mocks."""

    def __init__(self, regex):
        self._regex = re.compile(regex)

    def __eq__(self, string):
        return self._regex.search(string) is not None


class ImageDownloaderTestcast(unittest.TestCase):
    """The image downloader test class."""

    def __init__(self, *args, **kwargs):
        super(ImageDownloaderTestcast, self).__init__(*args, **kwargs)
        self.called_download_image = False
        self.called_uncompress_image = False
        self.called_get_build_id = False
        self.called_download_autotest_files = False
        self.called_download_debug_file = False

    @mock.patch.object(os, "makedirs")
    @mock.patch.object(os.path, "exists")
    def test_download_image(self, mock_path_exists, mock_mkdirs):
        # Set mock and test values.
        mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter)
        test_chroot = "/usr/local/home/chromeos"
        test_build_id = "lumpy-release/R36-5814.0.0"
        image_path = (
            "gs://chromeos-image-archive/%s/chromiumos_test_image.tar.xz"
            % test_build_id
        )

        downloader = download_images.ImageDownloader(
            logger_to_use=MOCK_LOGGER, cmd_exec=mock_cmd_exec
        )

        # Set os.path.exists to always return False and run downloader
        mock_path_exists.return_value = False
        test_flag.SetTestMode(True)
        self.assertRaises(
            download_images.MissingImage,
            downloader.DownloadImage,
            test_chroot,
            test_build_id,
            image_path,
        )

        # Verify os.path.exists was called thrice, with proper arguments.
        self.assertEqual(mock_path_exists.call_count, 3)
        mock_path_exists.assert_any_call(
            RegexMatcher(
                "/usr/local/home/chromeos/.*tmp/lumpy-release/"
                "R36-5814.0.0/chromiumos_test_image.bin"
            )
        )
        mock_path_exists.assert_any_call(
            RegexMatcher(
                "/usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0"
            )
        )
        mock_path_exists.assert_any_call("/etc/cros_chroot_version")

        # Verify we called os.mkdirs
        self.assertEqual(mock_mkdirs.call_count, 1)
        mock_mkdirs.assert_called_with(
            RegexMatcher(
                "/usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0"
            )
        )

        # Verify we called RunCommand once, with proper arguments.
        self.assertEqual(mock_cmd_exec.RunCommand.call_count, 1)
        expected_args = RegexMatcher(
            "/usr/local/home/chromeos/src/chromium/depot_tools/gsutil.py "
            "cp gs://chromeos-image-archive/lumpy-release/R36-5814.0.0/"
            "chromiumos_test_image.tar.xz "
            "/usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0"
        )

        mock_cmd_exec.RunCommand.assert_called_with(expected_args)

        # Reset the velues in the mocks; set os.path.exists to always return
        # True (except for "inside chroot" check).
        mock_path_exists.reset_mock()
        mock_cmd_exec.reset_mock()
        mock_path_exists.side_effect = lambda x: x != "/etc/cros_chroot_version"

        # Run downloader
        downloader.DownloadImage(test_chroot, test_build_id, image_path)

        # Verify os.path.exists was called thrice, with proper arguments.
        self.assertEqual(mock_path_exists.call_count, 3)
        mock_path_exists.assert_called_with(
            RegexMatcher(
                "/usr/local/home/chromeos/.*tmp/lumpy-release/"
                "R36-5814.0.0/chromiumos_test_image.bin"
            )
        )
        mock_path_exists.assert_any_call(
            RegexMatcher(
                "/usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0"
            )
        )
        mock_path_exists.assert_any_call("/etc/cros_chroot_version")

        # Verify we made no RunCommand or ChrootRunCommand calls (since
        # os.path.exists returned True, there was no work do be done).
        self.assertEqual(mock_cmd_exec.RunCommand.call_count, 0)
        self.assertEqual(mock_cmd_exec.ChrootRunCommand.call_count, 0)

    @mock.patch.object(os.path, "exists")
    def test_uncompress_image(self, mock_path_exists):
        # set mock and test values.
        mock_cmd_exec = mock.Mock(spec=command_executer.CommandExecuter)
        test_chroot = "/usr/local/home/chromeos"
        test_build_id = "lumpy-release/R36-5814.0.0"

        downloader = download_images.ImageDownloader(
            logger_to_use=MOCK_LOGGER, cmd_exec=mock_cmd_exec
        )

        # Set os.path.exists to always return False and run uncompress.
        mock_path_exists.return_value = False
        self.assertRaises(
            download_images.MissingImage,
            downloader.UncompressImage,
            test_chroot,
            test_build_id,
        )

        # Verify os.path.exists was called twice, with correct arguments.
        self.assertEqual(mock_path_exists.call_count, 2)
        mock_path_exists.assert_called_with(
            RegexMatcher(
                "/usr/local/home/chromeos/.*tmp/lumpy-release/"
                "R36-5814.0.0/chromiumos_test_image.bin"
            )
        )
        mock_path_exists.assert_any_call("/etc/cros_chroot_version")

        # Verify RunCommand was called twice with correct arguments.
        self.assertEqual(mock_cmd_exec.RunCommand.call_count, 2)
        # Call 1, should have 2 arguments
        self.assertEqual(len(mock_cmd_exec.RunCommand.call_args_list[0]), 2)
        actual_arg = mock_cmd_exec.RunCommand.call_args_list[0][0]
        expected_arg = (
            RegexMatcher(
                "cd /usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0 ; "
                "tar -Jxf chromiumos_test_image.tar.xz "
            ),
        )
        self.assertEqual(expected_arg, actual_arg)
        # 2nd arg must be exception handler
        except_handler_string = "RunCommandExceptionHandler.HandleException"
        self.assertTrue(
            except_handler_string
            in repr(mock_cmd_exec.RunCommand.call_args_list[0][1])
        )

        # Call 2, should have 2 arguments
        self.assertEqual(len(mock_cmd_exec.RunCommand.call_args_list[1]), 2)
        actual_arg = mock_cmd_exec.RunCommand.call_args_list[1][0]
        expected_arg = (
            RegexMatcher(
                "cd /usr/local/home/chromeos/.*tmp/lumpy-release/R36-5814.0.0 ; "
                "rm -f chromiumos_test_image.bin "
            ),
        )
        self.assertEqual(expected_arg, actual_arg)
        # 2nd arg must be empty
        self.assertTrue(
            "{}" in repr(mock_cmd_exec.RunCommand.call_args_list[1][1])
        )

        # Set os.path.exists to always return True (except for "inside chroot"
        # check) and run uncompress.
        mock_path_exists.reset_mock()
        mock_cmd_exec.reset_mock()
        mock_path_exists.side_effect = lambda x: x != "/etc/cros_chroot_version"
        downloader.UncompressImage(test_chroot, test_build_id)

        # Verify os.path.exists was called once, with correct arguments.
        self.assertEqual(mock_path_exists.call_count, 2)
        mock_path_exists.assert_called_with(
            RegexMatcher(
                "/usr/local/home/chromeos/.*tmp/lumpy-release/"
                "R36-5814.0.0/chromiumos_test_image.bin"
            )
        )
        mock_path_exists.assert_any_call("/etc/cros_chroot_version")

        # Verify RunCommand was not called.
        self.assertEqual(mock_cmd_exec.RunCommand.call_count, 0)

    def test_run(self):
        # Set test arguments
        test_chroot = "/usr/local/home/chromeos"
        test_build_id = "remote/lumpy/latest-dev"
        test_empty_autotest_path = ""
        test_empty_debug_path = ""
        test_autotest_path = "/tmp/autotest"
        test_debug_path = "/tmp/debug"
        download_debug = True

        # Set values to test/check.
        self.called_download_image = False
        self.called_uncompress_image = False
        self.called_get_build_id = False
        self.called_download_autotest_files = False
        self.called_download_debug_file = False

        # Define fake stub functions for Run to call
        def FakeGetBuildID(unused_root, unused_xbuddy_label):
            self.called_get_build_id = True
            return "lumpy-release/R36-5814.0.0"

        def GoodDownloadImage(root, build_id, image_path):
            if root or build_id or image_path:
                pass
            self.called_download_image = True
            return "chromiumos_test_image.bin"

        def BadDownloadImage(root, build_id, image_path):
            if root or build_id or image_path:
                pass
            self.called_download_image = True
            raise download_images.MissingImage("Could not download image")

        def FakeUncompressImage(root, build_id):
            if root or build_id:
                pass
            self.called_uncompress_image = True
            return 0

        def FakeDownloadAutotestFiles(root, build_id):
            if root or build_id:
                pass
            self.called_download_autotest_files = True
            return "autotest"

        def FakeDownloadDebugFile(root, build_id):
            if root or build_id:
                pass
            self.called_download_debug_file = True
            return "debug"

        # Initialize downloader
        downloader = download_images.ImageDownloader(logger_to_use=MOCK_LOGGER)

        # Set downloader to call fake stubs.
        downloader.GetBuildID = FakeGetBuildID
        downloader.UncompressImage = FakeUncompressImage
        downloader.DownloadImage = GoodDownloadImage
        downloader.DownloadAutotestFiles = FakeDownloadAutotestFiles
        downloader.DownloadDebugFile = FakeDownloadDebugFile

        # Call Run.
        image_path, autotest_path, debug_path = downloader.Run(
            test_chroot,
            test_build_id,
            test_empty_autotest_path,
            test_empty_debug_path,
            download_debug,
        )

        # Make sure it called both _DownloadImage and _UncompressImage
        self.assertTrue(self.called_download_image)
        self.assertTrue(self.called_uncompress_image)
        # Make sure it called DownloadAutotestFiles
        self.assertTrue(self.called_download_autotest_files)
        # Make sure it called DownloadDebugFile
        self.assertTrue(self.called_download_debug_file)
        # Make sure it returned an image and autotest path returned from this call
        self.assertTrue(image_path == "chromiumos_test_image.bin")
        self.assertTrue(autotest_path == "autotest")
        self.assertTrue(debug_path == "debug")

        # Call Run with a non-empty autotest and debug path
        self.called_download_autotest_files = False
        self.called_download_debug_file = False

        image_path, autotest_path, debug_path = downloader.Run(
            test_chroot,
            test_build_id,
            test_autotest_path,
            test_debug_path,
            download_debug,
        )

        # Verify that downloadAutotestFiles was not called
        self.assertFalse(self.called_download_autotest_files)
        # Make sure it returned the specified autotest path returned from this call
        self.assertTrue(autotest_path == test_autotest_path)
        # Make sure it returned the specified debug path returned from this call
        self.assertTrue(debug_path == test_debug_path)

        # Reset values; Now use fake stub that simulates DownloadImage failing.
        self.called_download_image = False
        self.called_uncompress_image = False
        self.called_download_autotest_files = False
        self.called_download_debug_file = False
        downloader.DownloadImage = BadDownloadImage

        # Call Run again.
        self.assertRaises(
            download_images.MissingImage,
            downloader.Run,
            test_chroot,
            test_autotest_path,
            test_debug_path,
            test_build_id,
            download_debug,
        )

        # Verify that UncompressImage and downloadAutotestFiles were not called,
        # since _DownloadImage "failed"
        self.assertTrue(self.called_download_image)
        self.assertFalse(self.called_uncompress_image)
        self.assertFalse(self.called_download_autotest_files)
        self.assertFalse(self.called_download_debug_file)


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