diff --git a/tools/releasetools/common.py b/tools/releasetools/common.py
index 632c1e2..171c794 100644
--- a/tools/releasetools/common.py
+++ b/tools/releasetools/common.py
@@ -834,7 +834,7 @@
   reconstructed from the block list, for whom we should avoid applying imgdiff.
 
   Args:
-    which: The partition name, which must be "system" or "vendor".
+    which: The partition name, e.g. "system", "vendor".
     tmpdir: The directory that contains the prebuilt image and block map file.
     input_zip: The target-files ZIP archive.
     allow_shared_blocks: Whether having shared blocks is allowed.
@@ -843,8 +843,6 @@
   Returns:
     A SparseImage object, with file_map info loaded.
   """
-  assert which in ("system", "vendor")
-
   path = os.path.join(tmpdir, "IMAGES", which + ".img")
   mappath = os.path.join(tmpdir, "IMAGES", which + ".map")
 
diff --git a/tools/releasetools/test_common.py b/tools/releasetools/test_common.py
index d4fa5f3..dce27fe 100644
--- a/tools/releasetools/test_common.py
+++ b/tools/releasetools/test_common.py
@@ -641,11 +641,13 @@
         },
         sparse_image.file_map)
 
-  def test_GetSparseImage_invalidImageName(self):
+  def test_GetSparseImage_missingImageFile(self):
     self.assertRaises(
-        AssertionError, common.GetSparseImage, 'system2', None, None, False)
+        AssertionError, common.GetSparseImage, 'system2', self.testdata_dir,
+        None, False)
     self.assertRaises(
-        AssertionError, common.GetSparseImage, 'unknown', None, None, False)
+        AssertionError, common.GetSparseImage, 'unknown', self.testdata_dir,
+        None, False)
 
   def test_GetSparseImage_missingBlockMapFile(self):
     target_files = common.MakeTempFile(prefix='target_files-', suffix='.zip')