diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index 8abc442..b7e114d 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -425,6 +425,52 @@
   return base_fs_file if exit_code == 0 else None
 
 
+def SetUpInDirAndFsConfig(origin_in, prop_dict):
+  """Returns the in_dir and fs_config that should be used for image building.
+
+  If the target uses system_root_image and it's building system.img, it creates
+  and returns a staged dir that combines the contents of /system (i.e. in the
+  given in_dir) and root.
+
+  Args:
+    origin_in: Path to the input directory.
+    prop_dict: A property dict that contains info like partition size. Values
+        may be updated.
+
+  Returns:
+    A tuple of in_dir and fs_config that should be used to build the image.
+  """
+  fs_config = prop_dict.get("fs_config")
+  if (prop_dict.get("system_root_image") != "true" or
+      prop_dict["mount_point"] != "system"):
+    return origin_in, fs_config
+
+  # Construct a staging directory of the root file system.
+  in_dir = common.MakeTempDir()
+  root_dir = prop_dict.get("root_dir")
+  if root_dir:
+    shutil.rmtree(in_dir)
+    shutil.copytree(root_dir, in_dir, symlinks=True)
+  in_dir_system = os.path.join(in_dir, "system")
+  shutil.rmtree(in_dir_system, ignore_errors=True)
+  shutil.copytree(origin_in, in_dir_system, symlinks=True)
+
+  # Change the mount point to "/".
+  prop_dict["mount_point"] = "/"
+  if fs_config:
+    # We need to merge the fs_config files of system and root.
+    merged_fs_config = common.MakeTempFile(
+        prefix="merged_fs_config", suffix=".txt")
+    with open(merged_fs_config, "w") as fw:
+      if "root_fs_config" in prop_dict:
+        with open(prop_dict["root_fs_config"]) as fr:
+          fw.writelines(fr.readlines())
+      with open(fs_config) as fr:
+        fw.writelines(fr.readlines())
+    fs_config = merged_fs_config
+  return in_dir, fs_config
+
+
 def CheckHeadroom(ext4fs_output, prop_dict):
   """Checks if there's enough headroom space available.
 
@@ -468,40 +514,25 @@
 
 
 def BuildImage(in_dir, prop_dict, out_file, target_out=None):
-  """Build an image to out_file from in_dir with property prop_dict.
-  After the function call, values in prop_dict is updated with
-  computed values.
+  """Builds an image for the files under in_dir and writes it to out_file.
+
+  When using system_root_image, it will additionally look for the files under
+  root (specified by 'root_dir') and builds an image that contains both sources.
 
   Args:
-    in_dir: path of input directory.
-    prop_dict: property dictionary.
-    out_file: path of the output image file.
-    target_out: path of the product out directory to read device specific FS
-        config files.
+    in_dir: Path to input directory.
+    prop_dict: A property dict that contains info like partition size. Values
+        will be updated with computed values.
+    out_file: The output image file.
+    target_out: Path to the TARGET_OUT directory as in Makefile. It actually
+        points to the /system directory under PRODUCT_OUT. fs_config (the one
+        under system/core/libcutils) reads device specific FS config files from
+        there.
 
   Returns:
     True iff the image is built successfully.
   """
-  # system_root_image=true: build a system.img that combines the contents of
-  # /system and root, which should be mounted at the root of the file system.
-  origin_in = in_dir
-  fs_config = prop_dict.get("fs_config")
-  if (prop_dict.get("system_root_image") == "true" and
-      prop_dict["mount_point"] == "system"):
-    in_dir = common.MakeTempDir()
-    # Change the mount point to "/".
-    prop_dict["mount_point"] = "/"
-    if fs_config:
-      # We need to merge the fs_config files of system and root.
-      merged_fs_config = common.MakeTempFile(prefix="merged_fs_config",
-                                             suffix=".txt")
-      with open(merged_fs_config, "w") as fw:
-        if "root_fs_config" in prop_dict:
-          with open(prop_dict["root_fs_config"]) as fr:
-            fw.writelines(fr.readlines())
-        with open(fs_config) as fr:
-          fw.writelines(fr.readlines())
-      fs_config = merged_fs_config
+  in_dir, fs_config = SetUpInDirAndFsConfig(in_dir, prop_dict)
 
   build_command = []
   fs_type = prop_dict.get("fs_type", "")
@@ -518,11 +549,11 @@
   if (prop_dict.get("use_logical_partitions") == "true" and
       "partition_size" not in prop_dict):
     # if partition_size is not defined, use output of `du' + reserved_size
-    success, size = GetDiskUsage(origin_in)
+    success, size = GetDiskUsage(in_dir)
     if not success:
       return False
     if OPTIONS.verbose:
-      print("The tree size of %s is %d MB." % (origin_in, size // BYTES_IN_MB))
+      print("The tree size of %s is %d MB." % (in_dir, size // BYTES_IN_MB))
     size += int(prop_dict.get("partition_reserved_size", 0))
     # Round this up to a multiple of 4K so that avbtool works
     size = common.RoundUpTo4K(size)
@@ -643,27 +674,17 @@
     print("Error: unknown filesystem type '%s'" % (fs_type))
     return False
 
-  if in_dir != origin_in:
-    # Construct a staging directory of the root file system.
-    root_dir = prop_dict.get("root_dir")
-    if root_dir:
-      shutil.rmtree(in_dir)
-      shutil.copytree(root_dir, in_dir, symlinks=True)
-    staging_system = os.path.join(in_dir, "system")
-    shutil.rmtree(staging_system, ignore_errors=True)
-    shutil.copytree(origin_in, staging_system, symlinks=True)
-
   (mkfs_output, exit_code) = RunCommand(build_command)
   if exit_code != 0:
     print("Error: '%s' failed with exit code %d:\n%s" % (
         build_command, exit_code, mkfs_output))
-    success, du = GetDiskUsage(origin_in)
+    success, du = GetDiskUsage(in_dir)
     du_str = ("%d bytes (%d MB)" % (du, du // BYTES_IN_MB)
              ) if success else "unknown"
     print(
         "Out of space? The tree size of {} is {}, with reserved space of {} "
         "bytes ({} MB).".format(
-            origin_in, du_str,
+            in_dir, du_str,
             int(prop_dict.get("partition_reserved_size", 0)),
             int(prop_dict.get("partition_reserved_size", 0)) // BYTES_IN_MB))
     if "original_partition_size" in prop_dict:
diff --git a/tools/releasetools/test_build_image.py b/tools/releasetools/test_build_image.py
index 161faff..19b5e08 100644
--- a/tools/releasetools/test_build_image.py
+++ b/tools/releasetools/test_build_image.py
@@ -14,10 +14,12 @@
 # limitations under the License.
 #
 
+import filecmp
+import os.path
 import unittest
 
 import common
-from build_image import CheckHeadroom, RunCommand
+from build_image import CheckHeadroom, RunCommand, SetUpInDirAndFsConfig
 
 
 class BuildImageTest(unittest.TestCase):
@@ -26,6 +28,9 @@
   EXT4FS_OUTPUT = (
       "Created filesystem with 2777/129024 inodes and 515099/516099 blocks")
 
+  def tearDown(self):
+    common.Cleanup()
+
   def test_CheckHeadroom_SizeUnderLimit(self):
     # Required headroom: 1000 blocks.
     prop_dict = {
@@ -91,4 +96,95 @@
     }
     self.assertFalse(CheckHeadroom(ext4fs_output, prop_dict))
 
-    common.Cleanup()
+  def test_SetUpInDirAndFsConfig_SystemRootImageFalse(self):
+    prop_dict = {
+        'fs_config': 'fs-config',
+        'mount_point': 'system',
+    }
+    in_dir, fs_config = SetUpInDirAndFsConfig('/path/to/in_dir', prop_dict)
+    self.assertEqual('/path/to/in_dir', in_dir)
+    self.assertEqual('fs-config', fs_config)
+    self.assertEqual('system', prop_dict['mount_point'])
+
+  def test_SetUpInDirAndFsConfig_SystemRootImageTrue_NonSystem(self):
+    prop_dict = {
+        'fs_config': 'fs-config',
+        'mount_point': 'vendor',
+        'system_root_image': 'true',
+    }
+    in_dir, fs_config = SetUpInDirAndFsConfig('/path/to/in_dir', prop_dict)
+    self.assertEqual('/path/to/in_dir', in_dir)
+    self.assertEqual('fs-config', fs_config)
+    self.assertEqual('vendor', prop_dict['mount_point'])
+
+  @staticmethod
+  def _gen_fs_config(partition):
+    fs_config = common.MakeTempFile(suffix='.txt')
+    with open(fs_config, 'w') as fs_config_fp:
+      fs_config_fp.write('fs-config-{}\n'.format(partition))
+    return fs_config
+
+  def test_SetUpInDirAndFsConfig_SystemRootImageTrue(self):
+    root_dir = common.MakeTempDir()
+    with open(os.path.join(root_dir, 'init'), 'w') as init_fp:
+      init_fp.write('init')
+
+    origin_in = common.MakeTempDir()
+    with open(os.path.join(origin_in, 'file'), 'w') as in_fp:
+      in_fp.write('system-file')
+    os.symlink('../etc', os.path.join(origin_in, 'symlink'))
+
+    fs_config_system = self._gen_fs_config('system')
+
+    prop_dict = {
+        'fs_config': fs_config_system,
+        'mount_point': 'system',
+        'root_dir': root_dir,
+        'system_root_image': 'true',
+    }
+    in_dir, fs_config = SetUpInDirAndFsConfig(origin_in, prop_dict)
+
+    self.assertTrue(filecmp.cmp(
+        os.path.join(in_dir, 'init'), os.path.join(root_dir, 'init')))
+    self.assertTrue(filecmp.cmp(
+        os.path.join(in_dir, 'system', 'file'),
+        os.path.join(origin_in, 'file')))
+    self.assertTrue(os.path.islink(os.path.join(in_dir, 'system', 'symlink')))
+
+    self.assertTrue(filecmp.cmp(fs_config_system, fs_config))
+    self.assertEqual('/', prop_dict['mount_point'])
+
+  def test_SetUpInDirAndFsConfig_SystemRootImageTrue_WithRootFsConfig(self):
+    root_dir = common.MakeTempDir()
+    with open(os.path.join(root_dir, 'init'), 'w') as init_fp:
+      init_fp.write('init')
+
+    origin_in = common.MakeTempDir()
+    with open(os.path.join(origin_in, 'file'), 'w') as in_fp:
+      in_fp.write('system-file')
+    os.symlink('../etc', os.path.join(origin_in, 'symlink'))
+
+    fs_config_system = self._gen_fs_config('system')
+    fs_config_root = self._gen_fs_config('root')
+
+    prop_dict = {
+        'fs_config': fs_config_system,
+        'mount_point': 'system',
+        'root_dir': root_dir,
+        'root_fs_config': fs_config_root,
+        'system_root_image': 'true',
+    }
+    in_dir, fs_config = SetUpInDirAndFsConfig(origin_in, prop_dict)
+
+    self.assertTrue(filecmp.cmp(
+        os.path.join(in_dir, 'init'), os.path.join(root_dir, 'init')))
+    self.assertTrue(filecmp.cmp(
+        os.path.join(in_dir, 'system', 'file'),
+        os.path.join(origin_in, 'file')))
+    self.assertTrue(os.path.islink(os.path.join(in_dir, 'system', 'symlink')))
+
+    with open(fs_config) as fs_config_fp:
+      fs_config_data = fs_config_fp.readlines()
+    self.assertIn('fs-config-system\n', fs_config_data)
+    self.assertIn('fs-config-root\n', fs_config_data)
+    self.assertEqual('/', prop_dict['mount_point'])