| #!/usr/bin/python |
| |
| # Copyright 2016, The Android Open Source Project |
| # |
| # Licensed under the Apache License, Version 2.0 (the "License"); |
| # you may not use this file except in compliance with the License. |
| # You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, software |
| # distributed under the License is distributed on an "AS IS" BASIS, |
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| # See the License for the specific language governing permissions and |
| # limitations under the License. |
| |
| |
| """Unit-test for bpttool.""" |
| |
| |
| import imp |
| import sys |
| import tempfile |
| import unittest |
| |
| sys.dont_write_bytecode = True |
| bpttool = imp.load_source('bpttool', './bpttool') |
| |
| |
| class FakeGuidGenerator(object): |
| """A GUID generator that dispenses predictable GUIDs. |
| |
| This is used in place of the usual GUID generator in bpttool which |
| is generating random GUIDs as per RFC 4122. |
| """ |
| |
| def dispense_guid(self, partition_number): |
| """Dispenses a new GUID.""" |
| |
| uuid = '01234567-89ab-cdef-0123-%012x' % partition_number |
| return uuid |
| |
| class PatternPartition(object): |
| """A partition image file containing a predictable pattern. |
| |
| This holds file data about a partition image file for binary pattern. |
| testing. |
| """ |
| def __init__(self, char='', file=None, partition_name=None, obj=None): |
| self.char = char |
| self.file = file |
| self.partition_name = partition_name |
| self.obj = obj |
| |
| class RoundToMultipleTest(unittest.TestCase): |
| """Unit tests for the RoundToMultiple() function.""" |
| |
| def testRoundUp(self): |
| """Checks that we round up correctly.""" |
| self.assertEqual(bpttool.RoundToMultiple(100, 10), 100) |
| self.assertEqual(bpttool.RoundToMultiple(189, 10), 190) |
| self.assertEqual(bpttool.RoundToMultiple(190, 10), 190) |
| self.assertEqual(bpttool.RoundToMultiple(191, 10), 200) |
| self.assertEqual(bpttool.RoundToMultiple(199, 10), 200) |
| self.assertEqual(bpttool.RoundToMultiple(200, 10), 200) |
| self.assertEqual(bpttool.RoundToMultiple(201, 10), 210) |
| self.assertEqual(bpttool.RoundToMultiple(-18, 10), -10) |
| self.assertEqual(bpttool.RoundToMultiple(-19, 10), -10) |
| self.assertEqual(bpttool.RoundToMultiple(-20, 10), -20) |
| self.assertEqual(bpttool.RoundToMultiple(-21, 10), -20) |
| |
| def testRoundDown(self): |
| """Checks that we round down correctly.""" |
| self.assertEqual(bpttool.RoundToMultiple(100, 10, True), 100) |
| self.assertEqual(bpttool.RoundToMultiple(189, 10, True), 180) |
| self.assertEqual(bpttool.RoundToMultiple(190, 10, True), 190) |
| self.assertEqual(bpttool.RoundToMultiple(191, 10, True), 190) |
| self.assertEqual(bpttool.RoundToMultiple(199, 10, True), 190) |
| self.assertEqual(bpttool.RoundToMultiple(200, 10, True), 200) |
| self.assertEqual(bpttool.RoundToMultiple(201, 10, True), 200) |
| self.assertEqual(bpttool.RoundToMultiple(-18, 10, True), -20) |
| self.assertEqual(bpttool.RoundToMultiple(-19, 10, True), -20) |
| self.assertEqual(bpttool.RoundToMultiple(-20, 10, True), -20) |
| self.assertEqual(bpttool.RoundToMultiple(-21, 10, True), -30) |
| |
| |
| class ParseSizeTest(unittest.TestCase): |
| """Unit tests for the ParseSize() function.""" |
| |
| def testIntegers(self): |
| """Checks parsing of integers.""" |
| self.assertEqual(bpttool.ParseSize(123), 123) |
| self.assertEqual(bpttool.ParseSize(17179869184), 1<<34) |
| self.assertEqual(bpttool.ParseSize(0x1000), 4096) |
| |
| def testRawNumbers(self): |
| """Checks parsing of raw numbers.""" |
| self.assertEqual(bpttool.ParseSize('123'), 123) |
| self.assertEqual(bpttool.ParseSize('17179869184'), 1<<34) |
| self.assertEqual(bpttool.ParseSize('0x1000'), 4096) |
| |
| def testDecimalUnits(self): |
| """Checks that decimal size units are interpreted correctly.""" |
| self.assertEqual(bpttool.ParseSize('1 kB'), pow(10, 3)) |
| self.assertEqual(bpttool.ParseSize('1 MB'), pow(10, 6)) |
| self.assertEqual(bpttool.ParseSize('1 GB'), pow(10, 9)) |
| self.assertEqual(bpttool.ParseSize('1 TB'), pow(10, 12)) |
| self.assertEqual(bpttool.ParseSize('1 PB'), pow(10, 15)) |
| |
| def testBinaryUnits(self): |
| """Checks that binary size units are interpreted correctly.""" |
| self.assertEqual(bpttool.ParseSize('1 KiB'), pow(2, 10)) |
| self.assertEqual(bpttool.ParseSize('1 MiB'), pow(2, 20)) |
| self.assertEqual(bpttool.ParseSize('1 GiB'), pow(2, 30)) |
| self.assertEqual(bpttool.ParseSize('1 TiB'), pow(2, 40)) |
| self.assertEqual(bpttool.ParseSize('1 PiB'), pow(2, 50)) |
| |
| def testFloatAndUnits(self): |
| """Checks that floating point numbers can be used with units.""" |
| self.assertEqual(bpttool.ParseSize('0.5 kB'), 500) |
| self.assertEqual(bpttool.ParseSize('0.5 KiB'), 512) |
| self.assertEqual(bpttool.ParseSize('0.5 GB'), 500000000) |
| self.assertEqual(bpttool.ParseSize('0.5 GiB'), 536870912) |
| self.assertEqual(bpttool.ParseSize('0.1 MiB'), 104858) |
| |
| class MakeDiskImageTest(unittest.TestCase): |
| """Unit tests for 'bpttool make_disk_image'.""" |
| |
| def setUp(self): |
| """Set-up method.""" |
| self.bpt = bpttool.Bpt() |
| |
| def _BinaryPattern(self, bpt_file_name, partition_patterns): |
| """Checks that a binary pattern may be written to a specified partition. |
| |
| This checks individual partion image writes to portions of a disk. Known |
| patterns are written into certain partitions and are verified after each |
| pattern has been written to. |
| |
| Arguments: |
| bpt_file_name: File name of bpt JSON containing partition information. |
| partition_patterns: List of tuples with each tuple having partition name |
| as the first argument, and character pattern as the |
| second argument. |
| |
| """ |
| bpt_file = open(bpt_file_name, 'r') |
| partitions_string, _ = self.bpt.make_table([bpt_file]) |
| bpt_tmp = tempfile.NamedTemporaryFile() |
| bpt_tmp.write(partitions_string) |
| bpt_tmp.seek(0) |
| partitions, _ = self.bpt._read_json([bpt_tmp]) |
| |
| # Declare list of partition images to be written and compared on disk. |
| pattern_images = [PatternPartition( |
| char=pp[1], |
| file=tempfile.NamedTemporaryFile(), |
| partition_name=pp[0]) |
| for pp in partition_patterns] |
| |
| # Store partition object and write a known character pattern image. |
| for pi in pattern_images: |
| pi.obj = [p for p in partitions if str(p.label) == pi.partition_name][0] |
| pi.file.write(bytearray(pi.char * int(pi.obj.size))) |
| |
| # Create the disk containing the partition filled with a known character |
| # pattern, seek to it's position and compare it to the supposed pattern. |
| with tempfile.NamedTemporaryFile() as generated_disk_image: |
| bpt_tmp.seek(0) |
| self.bpt.make_disk_image(generated_disk_image, |
| bpt_tmp, |
| [p.partition_name + ':' + p.file.name |
| for p in pattern_images]) |
| |
| for pi in pattern_images: |
| generated_disk_image.seek(pi.obj.offset) |
| pi.file.seek(0) |
| |
| self.assertEqual(generated_disk_image.read(pi.obj.size), |
| pi.file.read()) |
| pi.file.close() |
| |
| bpt_file.close() |
| bpt_tmp.close() |
| |
| def _LargeBinary(self, bpt_file_name): |
| """Helper function to write large partition images to disk images. |
| |
| This is a simple call to make_disk_image, passing a large in an image |
| which exceeds the it's size as specfied in the bpt file. |
| |
| Arguments: |
| bpt_file_name: File name of bpt JSON containing partition information. |
| |
| """ |
| with open(bpt_file_name, 'r') as bpt_file, \ |
| tempfile.NamedTemporaryFile() as bpt_tmp, \ |
| tempfile.NamedTemporaryFile() as generated_disk_image, \ |
| tempfile.NamedTemporaryFile() as large_partition_image: |
| partitions_string, _ = self.bpt.make_table([bpt_file]) |
| bpt_tmp.write(partitions_string) |
| bpt_tmp.seek(0) |
| partitions, _ = self.bpt._read_json([bpt_tmp]) |
| |
| # Create the over-sized partition image. |
| large_partition_image.write(bytearray('0' * |
| int(1.1*partitions[0].size + 1))) |
| |
| bpt_tmp.seek(0) |
| |
| # Expect exception here. |
| self.bpt.make_disk_image(generated_disk_image, bpt_tmp, |
| [p.label + ':' + large_partition_image.name for p in partitions]) |
| |
| def testBinaryPattern(self): |
| """Checks patterns written to partitions on disk images.""" |
| self._BinaryPattern('test/pattern_partition_single.bpt', [('charlie', 'c')]) |
| self._BinaryPattern('test/pattern_partition_multi.bpt', [('alpha', 'a'), |
| ('beta', 'b')]) |
| |
| def testExceedPartitionSize(self): |
| """Checks that exceedingly large partition images are not accepted.""" |
| try: |
| self._LargeBinary('test/pattern_partition_exceed_size.bpt') |
| except bpttool.BptError as e: |
| assert 'exceeds the partition size' in e.message |
| |
| |
| class MakeTableTest(unittest.TestCase): |
| """Unit tests for 'bpttool make_table'.""" |
| |
| def setUp(self): |
| """Reset GUID generator for every test.""" |
| self.fake_guid_generator = FakeGuidGenerator() |
| |
| def _MakeTable(self, input_file_names, |
| expected_json_file_name, |
| expected_gpt_file_name=None, |
| partitions_offset_begin=None, |
| disk_size=None, |
| disk_alignment=None, |
| disk_guid=None, |
| ab_suffixes=None): |
| """Helper function to create partition tables. |
| |
| This is a simple wrapper for Bpt.make_table() that compares its |
| output with an expected output when given a set of inputs. |
| |
| Arguments: |
| input_file_names: A list of .bpt files to process. |
| expected_json_file_name: File name of the file with expected JSON output. |
| expected_gpt_file_name: File name of the file with expected binary |
| output or None. |
| partitions_offset_begin: if not None, the size of the disk |
| partitions offset begin to use. |
| disk_size: if not None, the size of the disk to use. |
| disk_alignment: if not None, the disk alignment to use. |
| disk_guid: if not None, the disk GUID to use. |
| ab_suffixes: if not None, a list of the A/B suffixes (as a |
| comma-separated string) to use. |
| |
| """ |
| |
| inputs = [] |
| for file_name in input_file_names: |
| inputs.append(open(file_name)) |
| |
| bpt = bpttool.Bpt() |
| (json_str, gpt_bin) = bpt.make_table( |
| inputs, |
| partitions_offset_begin=partitions_offset_begin, |
| disk_size=disk_size, |
| disk_alignment=disk_alignment, |
| disk_guid=disk_guid, |
| ab_suffixes=ab_suffixes, |
| guid_generator=self.fake_guid_generator) |
| self.assertEqual(json_str, open(expected_json_file_name).read()) |
| |
| # Check that bpttool generates same JSON if being fed output it |
| # just generated without passing any options (e.g. disk size, |
| # alignment, suffixes). This verifies that we include all |
| # necessary information in the generated JSON. |
| (json_str2, _) = bpt.make_table( |
| [open(expected_json_file_name, 'r')], |
| guid_generator=self.fake_guid_generator) |
| self.assertEqual(json_str, json_str2) |
| |
| if expected_gpt_file_name: |
| self.assertEqual(gpt_bin, open(expected_gpt_file_name).read()) |
| |
| def testBase(self): |
| """Checks that the basic layout is as expected.""" |
| self._MakeTable(['test/base.bpt'], |
| 'test/expected_json_base.bpt', |
| disk_size=bpttool.ParseSize('10 GiB')) |
| |
| def testSize(self): |
| """Checks that disk size can be changed on the command-line.""" |
| self._MakeTable(['test/base.bpt'], |
| 'test/expected_json_size.bpt', |
| disk_size=bpttool.ParseSize('20 GiB')) |
| |
| def testAlignment(self): |
| """Checks that disk alignment can be changed on the command-line.""" |
| self._MakeTable(['test/base.bpt'], |
| 'test/expected_json_alignment.bpt', |
| disk_size=bpttool.ParseSize('10 GiB'), |
| disk_alignment=1048576) |
| |
| def testPartitionsOffsetBegin(self): |
| """Checks that disk partitions offset begin |
| can be changed on the command-line.""" |
| self._MakeTable(['test/base.bpt'], |
| 'test/expected_json_partitions_offset_begin.bpt', |
| partitions_offset_begin=bpttool.ParseSize('1 MiB'), |
| disk_size=bpttool.ParseSize('10 GiB')) |
| |
| def testDiskGuid(self): |
| """Checks that disk GUID can be changed on the command-line.""" |
| self._MakeTable(['test/base.bpt'], |
| 'test/expected_json_disk_guid.bpt', |
| disk_size=bpttool.ParseSize('10 GiB'), |
| disk_guid='01234567-89ab-cdef-0123-00000000002a') |
| |
| def testSuffixes(self): |
| """Checks that A/B-suffixes can be changed on the command-line.""" |
| self._MakeTable(['test/base.bpt'], |
| 'test/expected_json_suffixes.bpt', |
| disk_size=bpttool.ParseSize('10 GiB'), |
| ab_suffixes='-A,-B') |
| |
| def testStackedIgnore(self): |
| """Checks that partitions can be ignored through stacking.""" |
| self._MakeTable(['test/base.bpt', 'test/ignore_userdata.bpt'], |
| 'test/expected_json_stacked_ignore.bpt', |
| disk_size=bpttool.ParseSize('10 GiB')) |
| |
| def testStackedSize(self): |
| """Checks that partition size can be changed through stacking.""" |
| self._MakeTable(['test/base.bpt', 'test/change_odm_size.bpt'], |
| 'test/expected_json_stacked_size.bpt', |
| disk_size=bpttool.ParseSize('10 GiB')) |
| |
| def testStackedSettings(self): |
| """Checks that settings can be changed through stacking.""" |
| self._MakeTable(['test/base.bpt', 'test/override_settings.bpt'], |
| 'test/expected_json_stacked_override_settings.bpt') |
| |
| def testStackedNewPartition(self): |
| """Checks that a new partition can be added through stacking.""" |
| self._MakeTable(['test/base.bpt', 'test/new_partition.bpt'], |
| 'test/expected_json_stacked_new_partition.bpt', |
| disk_size=bpttool.ParseSize('10 GiB')) |
| |
| def testStackedChangeFlags(self): |
| """Checks that we can change partition flags through stacking.""" |
| self._MakeTable(['test/base.bpt', 'test/change_userdata_flags.bpt'], |
| 'test/expected_json_stacked_change_flags.bpt', |
| disk_size=bpttool.ParseSize('10 GiB')) |
| |
| def testStackedDisableAB(self): |
| """Checks that we can change disable A/B on partitions through stacking.""" |
| self._MakeTable(['test/base.bpt', 'test/disable_ab.bpt'], |
| 'test/expected_json_stacked_disable_ab.bpt', |
| disk_size=bpttool.ParseSize('10 GiB')) |
| |
| def testStackedNewPartitionOnTop(self): |
| """Checks that we can add a new partition only given the output JSON. |
| |
| This also tests that 'userdata' is shrunk (it has a size in the |
| input file 'expected_json_base.bpt') to make room for the new |
| partition. This works because the 'grow' attribute is preserved. |
| """ |
| self._MakeTable(['test/expected_json_base.bpt', |
| 'test/new_partition_on_top.bpt'], |
| 'test/expected_json_stacked_new_partition_on_top.bpt') |
| |
| def testStackedSizeAB(self): |
| """Checks that AB partition size can be changed given output JSON. |
| |
| This verifies that we un-expand A/B partitions when importing |
| output JSON. E.g. the output JSON has system_a, system_b which is |
| rewritten to just system as part of the import. |
| """ |
| self._MakeTable(['test/expected_json_base.bpt', |
| 'test/change_system_size.bpt'], |
| 'test/expected_json_stacked_change_ab_size.bpt') |
| |
| def testPositionAttribute(self): |
| """Checks that it's possible to influence partition order.""" |
| self._MakeTable(['test/base.bpt', |
| 'test/positions.bpt'], |
| 'test/expected_json_stacked_positions.bpt', |
| disk_size=bpttool.ParseSize('10 GiB')) |
| |
| def testBinaryOutput(self): |
| """Checks binary partition table output. |
| |
| This verifies that we generate the binary partition tables |
| correctly. |
| """ |
| self._MakeTable(['test/expected_json_stacked_change_flags.bpt'], |
| 'test/expected_json_stacked_change_flags.bpt', |
| 'test/expected_json_stacked_change_flags.bin') |
| |
| def testFileWithSyntaxErrors(self): |
| """Check that we catch errors in JSON files in a structured way.""" |
| try: |
| self._MakeTable(['test/file_with_syntax_errors.bpt'], |
| 'file_with_syntax_errors.bpt') |
| except bpttool.BptParsingError as e: |
| self.assertEqual(e.filename, 'test/file_with_syntax_errors.bpt') |
| |
| |
| class QueryPartitionTest(unittest.TestCase): |
| """Unit tests for 'bpttool query_partition'.""" |
| |
| def setUp(self): |
| """Set-up method.""" |
| self.bpt = bpttool.Bpt() |
| |
| def testQuerySize(self): |
| """Checks query for size.""" |
| self.assertEqual( |
| self.bpt.query_partition( |
| open('test/expected_json_stacked_change_flags.bpt'), |
| 'userdata', 'size', False), |
| '7449042944') |
| |
| def testQueryOffset(self): |
| """Checks query for offset.""" |
| self.assertEqual( |
| self.bpt.query_partition( |
| open('test/expected_json_stacked_change_flags.bpt'), |
| 'userdata', 'offset', False), |
| '3288354816') |
| |
| def testQueryGuid(self): |
| """Checks query for GUID.""" |
| self.assertEqual( |
| self.bpt.query_partition( |
| open('test/expected_json_stacked_change_flags.bpt'), |
| 'userdata', 'guid', False), |
| '01234567-89ab-cdef-0123-000000000007') |
| |
| def testQueryTypeGuid(self): |
| """Checks query for type GUID.""" |
| self.assertEqual( |
| self.bpt.query_partition( |
| open('test/expected_json_stacked_change_flags.bpt'), |
| 'userdata', 'type_guid', False), |
| '0bb7e6ed-4424-49c0-9372-7fbab465ab4c') |
| |
| def testQueryFlags(self): |
| """Checks query for flags.""" |
| self.assertEqual( |
| self.bpt.query_partition( |
| open('test/expected_json_stacked_change_flags.bpt'), |
| 'userdata', 'flags', False), |
| '0x0420000000000000') |
| |
| def testQuerySizeCollapse(self): |
| """Checks query for size when collapsing A/B partitions.""" |
| self.assertEqual( |
| self.bpt.query_partition( |
| open('test/expected_json_stacked_change_flags.bpt'), |
| 'odm', 'size', True), |
| '1073741824') |
| |
| |
| if __name__ == '__main__': |
| unittest.main() |
