1 files changed, 26 insertions, 2 deletions
diff --git a/tools/releasetools/test_utils.py b/tools/releasetools/test_utils.py
index e30d2b98e3..5bbcf7f355 100755
--- a/tools/releasetools/test_utils.py
+++ b/tools/releasetools/test_utils.py
@@ -19,6 +19,7 @@
 Utils for running unittests.
 """
 
+import avbtool
 import logging
 import os
 import os.path
@@ -57,12 +58,14 @@ def get_testdata_dir():
   current_dir = os.path.dirname(os.path.realpath(__file__))
   return os.path.join(current_dir, 'testdata')
 
+
 def get_current_dir():
   """Returns the current dir, relative to the script dir."""
   # The script dir is the one we want, which could be different from pwd.
   current_dir = os.path.dirname(os.path.realpath(__file__))
   return current_dir
 
+
 def get_search_path():
   """Returns the search path that has 'framework/signapk.jar' under."""
 
@@ -83,14 +86,33 @@ def get_search_path():
       # In relative to 'build/make/tools/releasetools' in the Android source.
       ['..'] * 4 + ['out', 'host', 'linux-x86'],
       # Or running the script unpacked from otatools.zip.
-      ['..']):
+          ['..']):
     full_path = os.path.realpath(os.path.join(current_dir, *path))
     if signapk_exists(full_path):
       return full_path
   return None
 
 
-def construct_sparse_image(chunks):
+def append_avb_footer(file_path: str, partition_name: str = ""):
+  avb = avbtool.AvbTool()
+  try:
+    args = ["avbtool", "add_hashtree_footer", "--image", file_path,
+            "--partition_name", partition_name, "--do_not_generate_fec"]
+    avb.run(args)
+  except SystemExit:
+    raise ValueError(f"Failed to append hashtree footer {args}")
+
+
+def erase_avb_footer(file_path: str):
+  avb = avbtool.AvbTool()
+  try:
+    args = ["avbtool", "erase_footer", "--image", file_path]
+    avb.run(args)
+  except SystemExit:
+    raise ValueError(f"Failed to erase hashtree footer {args}")
+
+
+def construct_sparse_image(chunks, partition_name: str = ""):
   """Returns a sparse image file constructed from the given chunks.
 
   From system/core/libsparse/sparse_format.h.
@@ -151,6 +173,7 @@ def construct_sparse_image(chunks):
       if data_size != 0:
         fp.write(os.urandom(data_size))
 
+  append_avb_footer(sparse_image, partition_name)
   return sparse_image
 
 
@@ -201,6 +224,7 @@ class ReleaseToolsTestCase(unittest.TestCase):
   def tearDown(self):
     common.Cleanup()
 
+
 class PropertyFilesTestCase(ReleaseToolsTestCase):
 
   @staticmethod