| # -*- coding: utf-8 -*- |
| # Copyright 2018 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. |
| |
| """Test the partition_lib module.""" |
| |
| from __future__ import print_function |
| |
| import os |
| |
| 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 import path_util |
| from chromite.lib.paygen import partition_lib |
| |
| |
| class PartitionLibTest(cros_test_lib.RunCommandTempDirTestCase): |
| """Test partition_lib functions.""" |
| |
| def testTruncate(self): |
| """Test truncating on extraction.""" |
| root = os.path.join(self.tempdir, 'root.bin') |
| |
| content = '0123456789' |
| osutils.WriteFile(root, content) |
| self.PatchObject(partition_lib, 'ExtractPartition') |
| self.PatchObject(partition_lib, 'Ext2FileSystemSize', return_value=2) |
| |
| partition_lib.ExtractRoot(None, root, truncate=False) |
| self.assertEqual(osutils.ReadFile(root), content) |
| |
| partition_lib.ExtractRoot(None, root) |
| self.assertEqual(osutils.ReadFile(root), content[:2]) |
| |
| def testExt2FileSystemSize(self): |
| """Test getting filesystem size on a simple output.""" |
| block_size = 4096 |
| block_count = 123 |
| |
| self.PatchObject(cros_build_lib, 'run', |
| return_value=cros_build_lib.CommandResult(output=""" |
| Block size: %d |
| Other thing: 123456798 |
| Not an integer: cdsad132csda |
| Block count: %d |
| """ % (block_size, block_count))) |
| |
| size = partition_lib.Ext2FileSystemSize('/dev/null') |
| self.assertEqual(size, block_size * block_count) |
| |
| def testExtractPartition(self): |
| """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, 4, 'fs', 'PART-A', ''), |
| image_lib.PartitionInfo(2, 4, 8, 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 testIsSquashfsImage(self): |
| """Tests we correctly identify a SquashFS image.""" |
| # Tests correct arguments are passed and SquashFS image is correctly |
| # identified. |
| # Don't need to test the path util functionality, make it just give us the |
| # image path back. |
| self.PatchObject(path_util, 'ToChrootPath', side_effect=lambda x: x) |
| image = '/foo/image' |
| self.assertTrue(partition_lib.IsSquashfsImage(image)) |
| self.assertCommandCalled(['unsquashfs', '-s', image], enter_chroot=True, |
| stdout=True) |
| |
| # Tests failure to identify. |
| self.PatchObject(cros_build_lib, 'run', |
| side_effect=cros_build_lib.RunCommandError('error')) |
| self.assertFalse(partition_lib.IsSquashfsImage(image)) |
| |
| def testIsExt4Image(self): |
| """Tests we correctly identify an Ext4 image.""" |
| # Tests correct arguments are passed and Ext4 image is correctly identified. |
| # Don't need to test the path util functionality, make it just give us the |
| # image path back. |
| self.PatchObject(path_util, 'ToChrootPath', side_effect=lambda x: x) |
| image = '/foo/image' |
| self.assertTrue(partition_lib.IsExt4Image(image)) |
| self.assertCommandCalled(['sudo', '--', 'tune2fs', '-l', image], |
| enter_chroot=True, |
| stdout=True) |
| |
| # Tests failure to identify. |
| self.PatchObject(cros_build_lib, 'run', |
| side_effect=cros_build_lib.RunCommandError('error')) |
| self.assertFalse(partition_lib.IsExt4Image(image)) |
| |
| def testIsGptImage(self): |
| """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): |
| """Tests we correctly identify different image types.""" |
| image = '/foo/image' |
| is_gpt = self.PatchObject(partition_lib, 'IsGptImage') |
| is_squashfs = self.PatchObject(partition_lib, 'IsSquashfsImage') |
| is_ext4 = self.PatchObject(partition_lib, 'IsExt4Image') |
| |
| 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_ext4.return_value = True |
| self.assertEqual(partition_lib.LookupImageType(image), |
| partition_lib.DLC_IMAGE) |
| |
| is_ext4.return_value = False |
| self.assertIsNone(partition_lib.LookupImageType(image)) |