diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index 49d1f0a..13201ef 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -376,6 +376,40 @@
     return None
   return base_fs_file
 
+
+def CheckHeadroom(ext4fs_output, prop_dict):
+  """Checks if there's enough headroom space available.
+
+  Headroom is the reserved space on system image (via PRODUCT_SYSTEM_HEADROOM),
+  which is useful for devices with low disk space that have system image
+  variation between builds. The 'partition_headroom' in prop_dict is the size
+  in bytes, while the numbers in 'ext4fs_output' are for 4K-blocks.
+
+  Args:
+    ext4fs_output: The output string from mke2fs command.
+    prop_dict: The property dict.
+
+  Returns:
+    The check result.
+  """
+  ext4fs_stats = re.compile(
+      r'Created filesystem with .* (?P<used_blocks>[0-9]+)/'
+      r'(?P<total_blocks>[0-9]+) blocks')
+  m = ext4fs_stats.match(ext4fs_output.strip().split('\n')[-1])
+  used_blocks = int(m.groupdict().get('used_blocks'))
+  total_blocks = int(m.groupdict().get('total_blocks'))
+  headroom_blocks = int(prop_dict.get('partition_headroom')) / BLOCK_SIZE
+  adjusted_blocks = total_blocks - headroom_blocks
+  if used_blocks > adjusted_blocks:
+    mount_point = prop_dict.get("mount_point")
+    print("Error: Not enough room on %s (total: %d blocks, used: %d blocks, "
+          "headroom: %d blocks, available: %d blocks)" % (
+              mount_point, total_blocks, used_blocks, headroom_blocks,
+              adjusted_blocks))
+    return False
+  return True
+
+
 def BuildImage(in_dir, prop_dict, out_file, target_out=None):
   """Build an image to out_file from in_dir with property prop_dict.
 
@@ -532,7 +566,6 @@
     shutil.copytree(origin_in, staging_system, symlinks=True)
 
   ext4fs_output = None
-
   try:
     if fs_type.startswith("ext4"):
       (ext4fs_output, exit_code) = RunCommand(build_command)
@@ -550,24 +583,10 @@
     print("Error: '%s' failed with exit code %d" % (build_command, exit_code))
     return False
 
-  # Check if there's enough headroom space available. This is useful for devices
-  # with low disk space that have system image variation between builds.
+  # Check if there's enough headroom space available for ext4 image.
   if "partition_headroom" in prop_dict and fs_type.startswith("ext4"):
     assert ext4fs_output is not None
-    ext4fs_stats = re.compile(
-        r'Created filesystem with .* (?P<used_blocks>[0-9]+)/'
-        r'(?P<total_blocks>[0-9]+) blocks')
-    m = ext4fs_stats.match(ext4fs_output.strip().split('\n')[-1])
-    used_blocks = int(m.groupdict().get('used_blocks'))
-    total_blocks = int(m.groupdict().get('total_blocks'))
-    headroom_blocks = int(prop_dict.get('partition_headroom')) / BLOCK_SIZE
-    adjusted_blocks = total_blocks - headroom_blocks
-    if used_blocks > adjusted_blocks:
-      mount_point = prop_dict.get("mount_point")
-      print("Error: Not enough room on %s (total: %d blocks, used: %d blocks, "
-            "headroom: %d blocks, available: %d blocks)" % (
-                mount_point, total_blocks, used_blocks,
-                headroom_blocks, adjusted_blocks))
+    if not CheckHeadroom(ext4fs_output, prop_dict):
       return False
 
   if not fs_spans_partition:
diff --git a/tools/releasetools/test_build_image.py b/tools/releasetools/test_build_image.py
new file mode 100644
index 0000000..6566a5a
--- /dev/null
+++ b/tools/releasetools/test_build_image.py
@@ -0,0 +1,65 @@
+#
+# Copyright (C) 2017 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.
+#
+
+import shutil
+import tempfile
+import unittest
+
+from build_image import CheckHeadroom, RunCommand
+
+
+class BuildImageTest(unittest.TestCase):
+
+  def test_CheckHeadroom_SizeUnderLimit(self):
+    ext4fs_output = ("Created filesystem with 2777/129024 inodes and "
+                     "508140/516099 blocks")
+    prop_dict = {
+        'partition_headroom' : '4194304',
+        'mount_point' : 'system',
+    }
+    self.assertTrue(CheckHeadroom(ext4fs_output, prop_dict))
+
+  def test_CheckHeadroom_InsufficientHeadroom(self):
+    ext4fs_output = ("Created filesystem with 2777/129024 inodes and "
+                     "515099/516099 blocks")
+    prop_dict = {
+        'partition_headroom' : '4100096',
+        'mount_point' : 'system',
+    }
+    self.assertFalse(CheckHeadroom(ext4fs_output, prop_dict))
+
+  def test_CheckHeadroom_WithMke2fsOutput(self):
+    """Tests the result parsing from actual call to mke2fs."""
+    input_dir = tempfile.mkdtemp()
+    output_image = tempfile.NamedTemporaryFile(suffix='.img')
+    command = ['mkuserimg_mke2fs.sh', input_dir, output_image.name, 'ext4',
+               '/system', '409600', '-j', '0']
+    ext4fs_output, exit_code = RunCommand(command)
+    self.assertEqual(0, exit_code)
+
+    prop_dict = {
+        'partition_headroom' : '40960',
+        'mount_point' : 'system',
+    }
+    self.assertTrue(CheckHeadroom(ext4fs_output, prop_dict))
+
+    prop_dict = {
+        'partition_headroom' : '413696',
+        'mount_point' : 'system',
+    }
+    self.assertFalse(CheckHeadroom(ext4fs_output, prop_dict))
+
+    shutil.rmtree(input_dir)