releasetools: Add unittest for common.GetSparseImage().

Add construct_sparse_image() to test_utils.py, which is a util function
to create sparse images. The new tests also partially cover the recent
changes that add 'incomplete' and 'uses_shard_blocks' tags.

Test: python -m unittest test_common
Change-Id: Ia15f5c4ad12423691216ebbad2c28f95c8427d7e
diff --git a/tools/releasetools/test_common.py b/tools/releasetools/test_common.py
index 6da286c..619eda5 100644
--- a/tools/releasetools/test_common.py
+++ b/tools/releasetools/test_common.py
@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+
 import os
 import tempfile
 import time
@@ -23,6 +24,7 @@
 import common
 import test_utils
 import validate_target_files
+from rangelib import RangeSet
 
 
 KiB = 1024
@@ -489,6 +491,147 @@
     self.assertRaises(AssertionError, common.ExtractPublicKey, wrong_input)
 
 
+class CommonUtilsTest(unittest.TestCase):
+
+  def tearDown(self):
+    common.Cleanup()
+
+  def test_GetSparseImage_emptyBlockMapFile(self):
+    target_files = common.MakeTempFile(prefix='target_files-', suffix='.zip')
+    with zipfile.ZipFile(target_files, 'w') as target_files_zip:
+      target_files_zip.write(
+          test_utils.construct_sparse_image([
+              (0xCAC1, 6),
+              (0xCAC3, 3),
+              (0xCAC1, 4)]),
+          arcname='IMAGES/system.img')
+      target_files_zip.writestr('IMAGES/system.map', '')
+      target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 8))
+      target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
+
+    tempdir, input_zip = common.UnzipTemp(target_files)
+    sparse_image = common.GetSparseImage('system', tempdir, input_zip, False)
+    input_zip.close()
+
+    self.assertDictEqual(
+        {
+            '__COPY': RangeSet("0"),
+            '__NONZERO-0': RangeSet("1-5 9-12"),
+        },
+        sparse_image.file_map)
+
+  def test_GetSparseImage_invalidImageName(self):
+    self.assertRaises(
+        AssertionError, common.GetSparseImage, 'system2', None, None, False)
+    self.assertRaises(
+        AssertionError, common.GetSparseImage, 'unknown', None, None, False)
+
+  def test_GetSparseImage_missingBlockMapFile(self):
+    target_files = common.MakeTempFile(prefix='target_files-', suffix='.zip')
+    with zipfile.ZipFile(target_files, 'w') as target_files_zip:
+      target_files_zip.write(
+          test_utils.construct_sparse_image([
+              (0xCAC1, 6),
+              (0xCAC3, 3),
+              (0xCAC1, 4)]),
+          arcname='IMAGES/system.img')
+      target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 8))
+      target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
+
+    tempdir, input_zip = common.UnzipTemp(target_files)
+    self.assertRaises(
+        AssertionError, common.GetSparseImage, 'system', tempdir, input_zip,
+        False)
+    input_zip.close()
+
+  def test_GetSparseImage_sharedBlocks_notAllowed(self):
+    """Tests the case of having overlapping blocks but disallowed."""
+    target_files = common.MakeTempFile(prefix='target_files-', suffix='.zip')
+    with zipfile.ZipFile(target_files, 'w') as target_files_zip:
+      target_files_zip.write(
+          test_utils.construct_sparse_image([(0xCAC2, 16)]),
+          arcname='IMAGES/system.img')
+      # Block 10 is shared between two files.
+      target_files_zip.writestr(
+          'IMAGES/system.map',
+          '\n'.join([
+              '/system/file1 1-5 9-10',
+              '/system/file2 10-12']))
+      target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 7))
+      target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
+
+    tempdir, input_zip = common.UnzipTemp(target_files)
+    self.assertRaises(
+        AssertionError, common.GetSparseImage, 'system', tempdir, input_zip,
+        False)
+    input_zip.close()
+
+  def test_GetSparseImage_sharedBlocks_allowed(self):
+    """Tests the case for target using BOARD_EXT4_SHARE_DUP_BLOCKS := true."""
+    target_files = common.MakeTempFile(prefix='target_files-', suffix='.zip')
+    with zipfile.ZipFile(target_files, 'w') as target_files_zip:
+      # Construct an image with a care_map of "0-5 9-12".
+      target_files_zip.write(
+          test_utils.construct_sparse_image([(0xCAC2, 16)]),
+          arcname='IMAGES/system.img')
+      # Block 10 is shared between two files.
+      target_files_zip.writestr(
+          'IMAGES/system.map',
+          '\n'.join([
+              '/system/file1 1-5 9-10',
+              '/system/file2 10-12']))
+      target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 7))
+      target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
+
+    tempdir, input_zip = common.UnzipTemp(target_files)
+    sparse_image = common.GetSparseImage('system', tempdir, input_zip, True)
+    input_zip.close()
+
+    self.assertDictEqual(
+        {
+            '__COPY': RangeSet("0"),
+            '__NONZERO-0': RangeSet("6-8 13-15"),
+            '/system/file1': RangeSet("1-5 9-10"),
+            '/system/file2': RangeSet("11-12"),
+        },
+        sparse_image.file_map)
+
+    # '/system/file2' should be marked with 'uses_shared_blocks', but not with
+    # 'incomplete'.
+    self.assertTrue(
+        sparse_image.file_map['/system/file2'].extra['uses_shared_blocks'])
+    self.assertNotIn(
+        'incomplete', sparse_image.file_map['/system/file2'].extra)
+
+    # All other entries should look normal without any tags.
+    self.assertFalse(sparse_image.file_map['__COPY'].extra)
+    self.assertFalse(sparse_image.file_map['__NONZERO-0'].extra)
+    self.assertFalse(sparse_image.file_map['/system/file1'].extra)
+
+  def test_GetSparseImage_incompleteRanges(self):
+    """Tests the case of ext4 images with holes."""
+    target_files = common.MakeTempFile(prefix='target_files-', suffix='.zip')
+    with zipfile.ZipFile(target_files, 'w') as target_files_zip:
+      target_files_zip.write(
+          test_utils.construct_sparse_image([(0xCAC2, 16)]),
+          arcname='IMAGES/system.img')
+      target_files_zip.writestr(
+          'IMAGES/system.map',
+          '\n'.join([
+              '/system/file1 1-5 9-10',
+              '/system/file2 11-12']))
+      target_files_zip.writestr('SYSTEM/file1', os.urandom(4096 * 7))
+      # '/system/file2' has less blocks listed (2) than actual (3).
+      target_files_zip.writestr('SYSTEM/file2', os.urandom(4096 * 3))
+
+    tempdir, input_zip = common.UnzipTemp(target_files)
+    sparse_image = common.GetSparseImage('system', tempdir, input_zip, False)
+    input_zip.close()
+
+    self.assertFalse(sparse_image.file_map['/system/file1'].extra)
+    self.assertTrue(sparse_image.file_map['/system/file2'].extra['incomplete'])
+
+
 class InstallRecoveryScriptFormatTest(unittest.TestCase):
   """Checks the format of install-recovery.sh.
 
diff --git a/tools/releasetools/test_utils.py b/tools/releasetools/test_utils.py
index ec53731..e64355b 100644
--- a/tools/releasetools/test_utils.py
+++ b/tools/releasetools/test_utils.py
@@ -18,7 +18,11 @@
 Utils for running unittests.
 """
 
+import os
 import os.path
+import struct
+
+import common
 
 
 def get_testdata_dir():
@@ -26,3 +30,67 @@
   # The script dir is the one we want, which could be different from pwd.
   current_dir = os.path.dirname(os.path.realpath(__file__))
   return os.path.join(current_dir, 'testdata')
+
+
+def construct_sparse_image(chunks):
+  """Returns a sparse image file constructed from the given chunks.
+
+  From system/core/libsparse/sparse_format.h.
+  typedef struct sparse_header {
+    __le32 magic;  // 0xed26ff3a
+    __le16 major_version;  // (0x1) - reject images with higher major versions
+    __le16 minor_version;  // (0x0) - allow images with higer minor versions
+    __le16 file_hdr_sz;  // 28 bytes for first revision of the file format
+    __le16 chunk_hdr_sz;  // 12 bytes for first revision of the file format
+    __le32 blk_sz;  // block size in bytes, must be a multiple of 4 (4096)
+    __le32 total_blks;  // total blocks in the non-sparse output image
+    __le32 total_chunks;  // total chunks in the sparse input image
+    __le32 image_checksum;  // CRC32 checksum of the original data, counting
+                            // "don't care" as 0. Standard 802.3 polynomial,
+                            // use a Public Domain table implementation
+  } sparse_header_t;
+
+  typedef struct chunk_header {
+    __le16 chunk_type;  // 0xCAC1 -> raw; 0xCAC2 -> fill;
+                        // 0xCAC3 -> don't care
+    __le16 reserved1;
+    __le32 chunk_sz;  // in blocks in output image
+    __le32 total_sz;  // in bytes of chunk input file including chunk header
+                      // and data
+  } chunk_header_t;
+
+  Args:
+    chunks: A list of chunks to be written. Each entry should be a tuple of
+        (chunk_type, block_number).
+
+  Returns:
+    Filename of the created sparse image.
+  """
+  SPARSE_HEADER_MAGIC = 0xED26FF3A
+  SPARSE_HEADER_FORMAT = "<I4H4I"
+  CHUNK_HEADER_FORMAT = "<2H2I"
+
+  sparse_image = common.MakeTempFile(prefix='sparse-', suffix='.img')
+  with open(sparse_image, 'wb') as fp:
+    fp.write(struct.pack(
+        SPARSE_HEADER_FORMAT, SPARSE_HEADER_MAGIC, 1, 0, 28, 12, 4096,
+        sum(chunk[1] for chunk in chunks),
+        len(chunks), 0))
+
+    for chunk in chunks:
+      data_size = 0
+      if chunk[0] == 0xCAC1:
+        data_size = 4096 * chunk[1]
+      elif chunk[0] == 0xCAC2:
+        data_size = 4
+      elif chunk[0] == 0xCAC3:
+        pass
+      else:
+        assert False, "Unsupported chunk type: {}".format(chunk[0])
+
+      fp.write(struct.pack(
+          CHUNK_HEADER_FORMAT, chunk[0], 0, chunk[1], data_size + 12))
+      if data_size != 0:
+        fp.write(os.urandom(data_size))
+
+  return sparse_image