blob: 3db05216aed7f98099dcc267865c9a06a20c10b3 [file] [log] [blame]
# 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.
"""Test the partition_lib module."""
import os
from chromite.lib import cgpt
from chromite.lib import cros_build_lib
from chromite.lib import cros_test_lib
from chromite.lib import image_lib
from chromite.lib import osutils
from chromite.lib.paygen import partition_lib
class PartitionLibTest(cros_test_lib.MockTempDirTestCase):
"""Test partition_lib functions with no mocks by default."""
def testTruncate(self) -> None:
"""Test truncating on extraction."""
root = self.tempdir / "root.bin"
# Create a small fs first.
osutils.AllocateFile(root, 1024 * 1024)
cros_build_lib.run(
["mke2fs", root],
extra_env={"PATH": "/sbin:/usr/sbin:%s" % os.environ["PATH"]},
)
# Then enlarge it.
os.truncate(root, 10 * 1024 * 1024)
self.PatchObject(partition_lib, "ExtractPartition")
partition_lib.ExtractRoot(None, root, truncate=False)
self.assertEqual(root.stat().st_size, 10 * 1024 * 1024)
partition_lib.ExtractRoot(None, root)
self.assertEqual(root.stat().st_size, 1024 * 1024)
def testExt2FileSystemSize(self) -> None:
"""Test getting filesystem size on a simple output."""
root = self.tempdir / "root.bin"
osutils.AllocateFile(root, 1024 * 1024)
cros_build_lib.run(
["mke2fs", root],
extra_env={"PATH": "/sbin:/usr/sbin:%s" % os.environ["PATH"]},
)
self.assertEqual(partition_lib.Ext2FileSystemSize(root), 1024 * 1024)
class PartitionLibMockTest(cros_test_lib.RunCommandTempDirTestCase):
"""Test partition_lib functions with run() mocked."""
def testExtractPartition(self) -> None:
"""Tests extraction on a simple image."""
part_a_bin = "0123"
part_b_bin = "4567"
image_bin = part_a_bin + part_b_bin
image = os.path.join(self.tempdir, "image.bin")
osutils.WriteFile(image, image_bin)
part_a = os.path.join(self.tempdir, "a.bin")
fake_partitions = (
image_lib.PartitionInfo(1, 0, 4, "fs", "PART-A"),
image_lib.PartitionInfo(2, 4, 4, "fs", "PART-B"),
)
self.PatchObject(
image_lib, "GetImageDiskPartitionInfo", return_value=fake_partitions
)
partition_lib.ExtractPartition(image, "PART-A", part_a)
self.assertEqual(osutils.ReadFile(part_a), part_a_bin)
def testIsGptImage(self) -> None:
"""Tests we correctly identify an Gpt image."""
# Tests correct arguments are passed and Gpt image is correctly
# identified.
image = "/foo/image"
part_info_mock = self.PatchObject(
image_lib, "GetImageDiskPartitionInfo"
)
self.assertTrue(partition_lib.IsGptImage(image))
part_info_mock.assert_called_once_with(image)
# Tests failure to identify.
part_info_mock.side_effect = cros_build_lib.RunCommandError("error")
part_info_mock.return_value = []
self.assertFalse(partition_lib.IsGptImage(image))
def testLookupImageType(self) -> None:
"""Tests we correctly identify different image types."""
image = "/foo/image"
is_gpt = self.PatchObject(partition_lib, "IsGptImage")
is_squashfs = self.PatchObject(image_lib, "IsSquashfsImage")
is_ext2 = self.PatchObject(image_lib, "IsExt2Image")
is_gpt.return_value = True
self.assertEqual(
partition_lib.LookupImageType(image), partition_lib.CROS_IMAGE
)
is_gpt.return_value = False
is_squashfs.return_value = True
self.assertEqual(
partition_lib.LookupImageType(image), partition_lib.DLC_IMAGE
)
is_squashfs.return_value = False
is_ext2.return_value = True
self.assertEqual(
partition_lib.LookupImageType(image), partition_lib.DLC_IMAGE
)
is_ext2.return_value = False
self.assertIsNone(partition_lib.LookupImageType(image))
def testHasMiniOSPartitions(self) -> None:
"""Tests we correctly identify miniOS supported images."""
image = "/foo/image"
self.PatchObject(cgpt.Disk, "FromImage", return_value=cgpt.Disk(""))
self.PatchObject(
cgpt.Disk, "GetPartitionByTypeGuid", side_effect=KeyError()
)
self.assertFalse(partition_lib.HasMiniOSPartitions(image))
self.PatchObject(cgpt.Disk, "GetPartitionByTypeGuid", return_value=[])
self.assertFalse(partition_lib.HasMiniOSPartitions(image))
self.PatchObject(
cgpt.Disk, "GetPartitionByTypeGuid", return_value=[None, None]
)
self.assertTrue(partition_lib.HasMiniOSPartitions(image))