diff --git a/delta_performer.cc b/delta_performer.cc
index 9392af3..28f61b8 100644
--- a/delta_performer.cc
+++ b/delta_performer.cc
@@ -60,15 +60,30 @@
 namespace {
 const int kUpdateStateOperationInvalid = -1;
 const int kMaxResumedUpdateFailures = 10;
+#if USE_MTD
+const int kUbiVolumeAttachTimeout = 5 * 60;
+#endif
 
 FileDescriptorPtr CreateFileDescriptor(const char* path) {
   FileDescriptorPtr ret;
 #if USE_MTD
-  if (UbiFileDescriptor::IsUbi(path)) {
-    ret.reset(new UbiFileDescriptor);
+  if (strstr(path, "/dev/ubi") == path) {
+    if (!UbiFileDescriptor::IsUbi(path)) {
+      // The volume might not have been attached at boot time.
+      int volume_no;
+      if (utils::SplitPartitionName(path, nullptr, &volume_no)) {
+        utils::TryAttachingUbiVolume(volume_no, kUbiVolumeAttachTimeout);
+      }
+    }
+    if (UbiFileDescriptor::IsUbi(path)) {
+      LOG(INFO) << path << " is a UBI device.";
+      ret.reset(new UbiFileDescriptor);
+    }
   } else if (MtdFileDescriptor::IsMtd(path)) {
+    LOG(INFO) << path << " is an MTD device.";
     ret.reset(new MtdFileDescriptor);
   } else {
+    LOG(INFO) << path << " is not an MTD nor a UBI device.";
 #endif
     ret.reset(new EintrSafeFileDescriptor);
 #if USE_MTD
@@ -81,8 +96,15 @@
 // and sets *err to 0. On failure, sets *err to errno and returns nullptr.
 FileDescriptorPtr OpenFile(const char* path, int* err) {
   FileDescriptorPtr fd = CreateFileDescriptor(path);
-  // TODO(namnguyen): If we're working with MTD or UBI, DO NOT use O_RDWR.
-  if (!fd->Open(path, O_RDWR, 000)) {
+  int mode = O_RDWR;
+#if USE_MTD
+  // On NAND devices, we can either read, or write, but not both. So here we
+  // use O_WRONLY.
+  if (UbiFileDescriptor::IsUbi(path) || MtdFileDescriptor::IsMtd(path)) {
+    mode = O_WRONLY;
+  }
+#endif
+  if (!fd->Open(path, mode, 000)) {
     *err = errno;
     PLOG(ERROR) << "Unable to open file " << path;
     return nullptr;
diff --git a/mtd_file_descriptor.cc b/mtd_file_descriptor.cc
index 1b33d05..2e0d329 100644
--- a/mtd_file_descriptor.cc
+++ b/mtd_file_descriptor.cc
@@ -10,24 +10,32 @@
 #include <sys/ioctl.h>
 #include <sys/stat.h>
 #include <sys/types.h>
+#include <vector>
 
 #include <base/files/file_path.h>
 #include <base/strings/string_number_conversions.h>
+#include <base/strings/string_util.h>
+#include <base/strings/stringprintf.h>
+#include <update_engine/subprocess.h>
 
 #include "update_engine/utils.h"
 
+using std::string;
+using std::vector;
+
 namespace {
 
 static const char kSysfsClassUbi[] = "/sys/class/ubi/";
 static const char kUsableEbSize[] = "/usable_eb_size";
 static const char kReservedEbs[] = "/reserved_ebs";
 
+using chromeos_update_engine::Subprocess;
 using chromeos_update_engine::UbiVolumeInfo;
 using chromeos_update_engine::utils::ReadFile;
 
 // Return a UbiVolumeInfo pointer if |path| is a UBI volume. Otherwise, return
 // a null unique pointer.
-std::unique_ptr<UbiVolumeInfo> GetUbiVolumeInfo(const char* path) {
+std::unique_ptr<UbiVolumeInfo> GetUbiVolumeInfo(const string& path) {
   base::FilePath device_node(path);
   base::FilePath ubi_name(device_node.BaseName());
 
@@ -39,23 +47,33 @@
   // Obtain volume info from sysfs.
   std::string s_reserved_ebs;
   if (!ReadFile(sysfs_node + kReservedEbs, &s_reserved_ebs)) {
+    LOG(ERROR) << "Cannot read " << sysfs_node + kReservedEbs;
     return ret;
   }
   std::string s_eb_size;
   if (!ReadFile(sysfs_node + kUsableEbSize, &s_eb_size)) {
+    LOG(ERROR) << "Cannot read " << sysfs_node + kUsableEbSize;
     return ret;
   }
 
-  size_t reserved_ebs, eb_size;
-  if (!base::StringToSizeT(s_reserved_ebs, &reserved_ebs)) {
+  base::TrimWhitespaceASCII(s_reserved_ebs,
+                            base::TRIM_TRAILING,
+                            &s_reserved_ebs);
+  base::TrimWhitespaceASCII(s_eb_size, base::TRIM_TRAILING, &s_eb_size);
+
+  uint64_t reserved_ebs, eb_size;
+  if (!base::StringToUint64(s_reserved_ebs, &reserved_ebs)) {
+    LOG(ERROR) << "Cannot parse reserved_ebs: " << s_reserved_ebs;
     return ret;
   }
-  if (!base::StringToSizeT(s_eb_size, &eb_size)) {
+  if (!base::StringToUint64(s_eb_size, &eb_size)) {
+    LOG(ERROR) << "Cannot parse usable_eb_size: " << s_eb_size;
     return ret;
   }
 
   ret.reset(new UbiVolumeInfo);
-  ret->size = reserved_ebs * eb_size;
+  ret->reserved_ebs = reserved_ebs;
+  ret->eraseblock_size = eb_size;
   return ret;
 }
 
@@ -74,14 +92,22 @@
 
 bool MtdFileDescriptor::Open(const char* path, int flags, mode_t mode) {
   // This File Descriptor does not support read and write.
-  TEST_AND_RETURN_FALSE((flags & O_RDWR) != O_RDWR);
+  TEST_AND_RETURN_FALSE((flags & O_ACCMODE) != O_RDWR);
+  // But we need to open the underlying file descriptor in O_RDWR mode because
+  // during write, we need to read back to verify the write actually sticks or
+  // we have to skip the block. That job is done by mtdutils library.
+  if ((flags & O_ACCMODE) == O_WRONLY) {
+    flags &= ~O_ACCMODE;
+    flags |= O_RDWR;
+  }
   TEST_AND_RETURN_FALSE(
       EintrSafeFileDescriptor::Open(path, flags | O_CLOEXEC, mode));
 
-  if (flags & O_RDONLY) {
-    read_ctx_.reset(mtd_read_descriptor(fd_, path));
-  } else if (flags & O_WRONLY) {
+  if ((flags & O_ACCMODE) == O_RDWR) {
     write_ctx_.reset(mtd_write_descriptor(fd_, path));
+    nr_written_ = 0;
+  } else {
+    read_ctx_.reset(mtd_read_descriptor(fd_, path));
   }
 
   if (!read_ctx_ && !write_ctx_) {
@@ -105,50 +131,66 @@
 
 ssize_t MtdFileDescriptor::Write(const void* buf, size_t count) {
   CHECK(write_ctx_);
-  return mtd_write_data(write_ctx_.get(), static_cast<const char*>(buf), count);
+  ssize_t result = mtd_write_data(write_ctx_.get(),
+                                  static_cast<const char*>(buf),
+                                  count);
+  if (result > 0) {
+    nr_written_ += result;
+  }
+  return result;
 }
 
 off64_t MtdFileDescriptor::Seek(off64_t offset, int whence) {
-  CHECK(read_ctx_);
+  if (write_ctx_) {
+    // Ignore seek in write mode.
+    return nr_written_;
+  }
   return EintrSafeFileDescriptor::Seek(offset, whence);
 }
 
-void MtdFileDescriptor::Reset() {
-  EintrSafeFileDescriptor::Reset();
+bool MtdFileDescriptor::Close() {
   read_ctx_.reset();
   write_ctx_.reset();
+  return EintrSafeFileDescriptor::Close();
 }
 
 bool UbiFileDescriptor::IsUbi(const char* path) {
+  base::FilePath device_node(path);
+  base::FilePath ubi_name(device_node.BaseName());
+  TEST_AND_RETURN_FALSE(StartsWithASCII(ubi_name.MaybeAsASCII(), "ubi", true));
+
   return static_cast<bool>(GetUbiVolumeInfo(path));
 }
 
-std::unique_ptr<UbiVolumeInfo> UbiFileDescriptor::CreateWriteContext(
-    const char* path) {
-  std::unique_ptr<UbiVolumeInfo> info = GetUbiVolumeInfo(path);
-  uint64_t volume_size;
-  if (info && (ioctl(fd_, UBI_IOCVOLUP, &volume_size) != 0 ||
-               volume_size != info->size)) {
-    info.reset();
-  }
-  return info;
-}
-
 bool UbiFileDescriptor::Open(const char* path, int flags, mode_t mode) {
+  std::unique_ptr<UbiVolumeInfo> info = GetUbiVolumeInfo(path);
+  if (!info) {
+    return false;
+  }
+
   // This File Descriptor does not support read and write.
-  TEST_AND_RETURN_FALSE((flags & O_RDWR) != O_RDWR);
+  TEST_AND_RETURN_FALSE((flags & O_ACCMODE) != O_RDWR);
   TEST_AND_RETURN_FALSE(
       EintrSafeFileDescriptor::Open(path, flags | O_CLOEXEC, mode));
 
-  if (flags & O_RDONLY) {
-    read_ctx_ = GetUbiVolumeInfo(path);
-  } else if (flags & O_WRONLY) {
-    write_ctx_ = CreateWriteContext(path);
-  }
+  usable_eb_blocks_ = info->reserved_ebs;
+  eraseblock_size_ = info->eraseblock_size;
+  volume_size_ = usable_eb_blocks_ * eraseblock_size_;
 
-  if (!read_ctx_ && !write_ctx_) {
-    Close();
-    return false;
+  if ((flags & O_ACCMODE) == O_WRONLY) {
+    // It's best to use volume update ioctl so that UBI layer will mark the
+    // volume as being updated, and only clear that mark if the update is
+    // successful. We will need to pad to the whole volume size at close.
+    uint64_t vsize = volume_size_;
+    if (ioctl(fd_, UBI_IOCVOLUP, &vsize) != 0) {
+      PLOG(ERROR) << "Cannot issue volume update ioctl";
+      EintrSafeFileDescriptor::Close();
+      return false;
+    }
+    mode_ = kWriteOnly;
+    nr_written_ = 0;
+  } else {
+    mode_ = kReadOnly;
   }
 
   return true;
@@ -161,24 +203,51 @@
 }
 
 ssize_t UbiFileDescriptor::Read(void* buf, size_t count) {
-  CHECK(read_ctx_);
+  CHECK(mode_ == kReadOnly);
   return EintrSafeFileDescriptor::Read(buf, count);
 }
 
 ssize_t UbiFileDescriptor::Write(const void* buf, size_t count) {
-  CHECK(write_ctx_);
-  return EintrSafeFileDescriptor::Write(buf, count);
+  CHECK(mode_ == kWriteOnly);
+  ssize_t nr_chunk = EintrSafeFileDescriptor::Write(buf, count);
+  if (nr_chunk >= 0) {
+    nr_written_ += nr_chunk;
+  }
+  return nr_chunk;
 }
 
 off64_t UbiFileDescriptor::Seek(off64_t offset, int whence) {
-  CHECK(read_ctx_);
+  if (mode_ == kWriteOnly) {
+    // Ignore seek in write mode.
+    return nr_written_;
+  }
   return EintrSafeFileDescriptor::Seek(offset, whence);
 }
 
-void UbiFileDescriptor::Reset() {
-  EintrSafeFileDescriptor::Reset();
-  read_ctx_.reset();
-  write_ctx_.reset();
+bool UbiFileDescriptor::Close() {
+  bool pad_ok = true;
+  if (IsOpen() && mode_ == kWriteOnly) {
+    char buf[1024];
+    memset(buf, 0xFF, sizeof(buf));
+    while (nr_written_ < volume_size_) {
+      // We have written less than the whole volume. In order for us to clear
+      // the update marker, we need to fill the rest. It is recommended to fill
+      // UBI writes with 0xFF.
+      uint64_t to_write = volume_size_ - nr_written_;
+      if (to_write > sizeof(buf)) {
+        to_write = sizeof(buf);
+      }
+      ssize_t nr_chunk = EintrSafeFileDescriptor::Write(buf, to_write);
+      if (nr_chunk < 0) {
+        LOG(ERROR) << "Cannot 0xFF-pad before closing.";
+        // There is an error, but we can't really do any meaningful thing here.
+        pad_ok = false;
+        break;
+      }
+      nr_written_ += nr_chunk;
+    }
+  }
+  return EintrSafeFileDescriptor::Close() && pad_ok;
 }
 
 }  // namespace chromeos_update_engine
diff --git a/mtd_file_descriptor.h b/mtd_file_descriptor.h
index 4694016..ad545b6 100644
--- a/mtd_file_descriptor.h
+++ b/mtd_file_descriptor.h
@@ -28,28 +28,28 @@
   ssize_t Read(void* buf, size_t count) override;
   ssize_t Write(const void* buf, size_t count) override;
   off64_t Seek(off64_t offset, int whence) override;
-  void Reset() override;
+  bool Close() override;
 
  private:
   std::unique_ptr<MtdReadContext, decltype(&mtd_read_close)> read_ctx_;
   std::unique_ptr<MtdWriteContext, decltype(&mtd_write_close)> write_ctx_;
+  uint64_t nr_written_;
 };
 
-// TODO(namnguyen) This is a placeholder struct. This struct, and the
-// UbiFileDescriptor class below will need finalized later.
 struct UbiVolumeInfo {
-  uint64_t size;
+  // Number of eraseblocks.
+  uint64_t reserved_ebs;
+  // Size of each eraseblock.
+  uint64_t eraseblock_size;
 };
 
 // A file descriptor to update a UBI volume, similar to MtdFileDescriptor.
 // Once the file descriptor is opened for write, the volume is marked as being
 // updated. The volume will not be usable until an update is completed. See
 // UBI_IOCVOLUP ioctl operation.
-// TODO(namnguyen) Again, this needs fleshed out when we have better library to
-// interact with UBI volumes. I would expect this class to be very similar to
-// MtdFileDescriptor, with two different contexts to bridge C-C++ divide.
 class UbiFileDescriptor : public EintrSafeFileDescriptor {
  public:
+  // Perform some queries about |path| to see if it is a UBI volume.
   static bool IsUbi(const char* path);
 
   bool Open(const char* path, int flags, mode_t mode) override;
@@ -57,13 +57,20 @@
   ssize_t Read(void* buf, size_t count) override;
   ssize_t Write(const void* buf, size_t count) override;
   off64_t Seek(off64_t offset, int whence) override;
-  void Reset() override;
+  bool Close() override;
 
  private:
-  std::unique_ptr<UbiVolumeInfo> CreateWriteContext(const char* path);
+  enum Mode {
+    kReadOnly,
+    kWriteOnly
+  };
 
-  std::unique_ptr<UbiVolumeInfo> read_ctx_;
-  std::unique_ptr<UbiVolumeInfo> write_ctx_;
+  uint64_t usable_eb_blocks_;
+  uint64_t eraseblock_size_;
+  uint64_t volume_size_;
+  uint64_t nr_written_;
+
+  Mode mode_;
 };
 
 }  // namespace chromeos_update_engine
diff --git a/postinstall_runner_action.cc b/postinstall_runner_action.cc
index b619cce..f27aace 100644
--- a/postinstall_runner_action.cc
+++ b/postinstall_runner_action.cc
@@ -38,8 +38,18 @@
                                            &temp_rootfs_dir_));
   ScopedDirRemover temp_dir_remover(temp_rootfs_dir_);
 
-  if (!utils::MountFilesystem(install_device, temp_rootfs_dir_, MS_RDONLY))
+  const string mountable_device =
+      utils::MakePartitionNameForMount(install_device);
+  if (mountable_device.empty()) {
+    LOG(ERROR) << "Cannot make mountable device from " << install_device;
     return;
+  }
+
+  if (!utils::MountFilesystem(mountable_device, temp_rootfs_dir_, MS_RDONLY))
+    return;
+
+  LOG(INFO) << "Performing postinst with install device " << install_device
+            << " and mountable device " << mountable_device;
 
   temp_dir_remover.set_should_remove(false);
   completer.set_should_complete(false);
diff --git a/utils.cc b/utils.cc
index 1f57672..7b3c6d0 100644
--- a/utils.cc
+++ b/utils.cc
@@ -71,6 +71,48 @@
 // in GetFileFormat.
 const int kGetFileFormatMaxHeaderSize = 32;
 
+// Return true if |disk_name| is an MTD or a UBI device. Note that this test is
+// simply based on the name of the device.
+bool IsMtdDeviceName(const string& disk_name) {
+  return StartsWithASCII(disk_name, "/dev/ubi", true) ||
+         StartsWithASCII(disk_name, "/dev/mtd", true);
+}
+
+// Return the device name for the corresponding partition on a NAND device.
+// WARNING: This function returns device names that are not mountable.
+string MakeNandPartitionName(int partition_num) {
+  switch (partition_num) {
+    case 2:
+    case 4:
+    case 6: {
+      return base::StringPrintf("/dev/mtd%d", partition_num);
+    }
+    default: {
+      return base::StringPrintf("/dev/ubi%d_0", partition_num);
+    }
+  }
+}
+
+// Return the device name for the corresponding partition on a NAND device that
+// may be mountable (but may not be writable).
+string MakeNandPartitionNameForMount(int partition_num) {
+  switch (partition_num) {
+    case 2:
+    case 4:
+    case 6: {
+      return base::StringPrintf("/dev/mtd%d", partition_num);
+    }
+    case 3:
+    case 5:
+    case 7: {
+      return base::StringPrintf("/dev/ubiblock%d_0", partition_num);
+    }
+    default: {
+      return base::StringPrintf("/dev/ubi%d_0", partition_num);
+    }
+  }
+}
+
 }  // namespace
 
 namespace utils {
@@ -107,11 +149,6 @@
   string disk_name;
   int partition_num;
   if (SplitPartitionName(boot_device, &disk_name, &partition_num)) {
-    if (disk_name == "/dev/ubiblock") {
-      // Special case for NAND devices.
-      // eg: /dev/ubiblock3_0 becomes /dev/mtdblock2
-      disk_name = "/dev/mtdblock";
-    }
     // Currently this assumes the partition number of the boot device is
     // 3, 5, or 7, and changes it to 2, 4, or 6, respectively, to
     // get the kernel device.
@@ -450,14 +487,21 @@
 }
 
 string MakePartitionName(const string& disk_name, int partition_num) {
+  if (partition_num < 1) {
+    LOG(ERROR) << "Invalid partition number: " << partition_num;
+    return string();
+  }
+
   if (!StartsWithASCII(disk_name, "/dev/", true)) {
     LOG(ERROR) << "Invalid disk name: " << disk_name;
     return string();
   }
 
-  if (partition_num < 1) {
-    LOG(ERROR) << "Invalid partition number: " << partition_num;
-    return string();
+  if (IsMtdDeviceName(disk_name)) {
+    // Special case for UBI block devices.
+    //   1. ubiblock is not writable, we need to use plain "ubi".
+    //   2. There is a "_0" suffix.
+    return MakeNandPartitionName(partition_num);
   }
 
   string partition_name = disk_name;
@@ -470,13 +514,20 @@
 
   partition_name += std::to_string(partition_num);
 
-  if (StartsWithASCII(partition_name, "/dev/ubiblock", true)) {
-    // Special case for UBI block devieces that have "_0" suffix.
-    partition_name += "_0";
-  }
   return partition_name;
 }
 
+string MakePartitionNameForMount(const string& part_name) {
+  if (IsMtdDeviceName(part_name)) {
+    int partition_num;
+    if (!SplitPartitionName(part_name, nullptr, &partition_num)) {
+      return "";
+    }
+    return MakeNandPartitionNameForMount(partition_num);
+  }
+  return part_name;
+}
+
 string SysfsBlockDevice(const string& device) {
   base::FilePath device_path(device);
   if (device_path.DirName().value() != "/dev") {
@@ -530,6 +581,39 @@
   return S_ISDIR(stbuf.st_mode);
 }
 
+bool TryAttachingUbiVolume(int volume_num, int timeout) {
+  const string volume_path = base::StringPrintf("/dev/ubi%d_0", volume_num);
+  if (FileExists(volume_path.c_str())) {
+    return true;
+  }
+
+  int exit_code;
+  vector<string> cmd = {
+      "ubiattach",
+      "-m",
+      base::StringPrintf("%d", volume_num),
+      "-d",
+      base::StringPrintf("%d", volume_num)
+  };
+  TEST_AND_RETURN_FALSE(Subprocess::SynchronousExec(cmd, &exit_code, nullptr));
+  TEST_AND_RETURN_FALSE(exit_code == 0);
+
+  cmd = {
+      "ubiblock",
+      "--create",
+      volume_path
+  };
+  TEST_AND_RETURN_FALSE(Subprocess::SynchronousExec(cmd, &exit_code, nullptr));
+  TEST_AND_RETURN_FALSE(exit_code == 0);
+
+  while (timeout > 0 && !FileExists(volume_path.c_str())) {
+    sleep(1);
+    timeout--;
+  }
+
+  return FileExists(volume_path.c_str());
+}
+
 // If |path| is absolute, or explicit relative to the current working directory,
 // leaves it as is. Otherwise, if TMPDIR is defined in the environment and is
 // non-empty, prepends it to |path|. Otherwise, prepends /tmp.  Returns the
diff --git a/utils.h b/utils.h
index 8e168f1..c820ab2 100644
--- a/utils.h
+++ b/utils.h
@@ -128,6 +128,11 @@
 // Returns true if |path| exists and is a directory.
 bool IsDir(const char* path);
 
+// Try attaching UBI |volume_num|. If there is any error executing required
+// commands to attach the volume, this function returns false. This function
+// only returns true if "/dev/ubi%d_0" becomes available in |timeout| seconds.
+bool TryAttachingUbiVolume(int volume_num, int timeout);
+
 // If |base_filename_template| is neither absolute (starts with "/") nor
 // explicitly relative to the current working directory (starts with "./" or
 // "../"), then it is prepended the value of TMPDIR, which defaults to /tmp if
@@ -179,6 +184,14 @@
 std::string MakePartitionName(const std::string& disk_name,
                               int partition_num);
 
+// Similar to "MakePartitionName" but returns a name that is suitable for
+// mounting. On NAND system we can write to "/dev/ubiX_0", which is what
+// MakePartitionName returns, but we cannot mount that device. To mount, we
+// have to use "/dev/ubiblockX_0" for rootfs. Stateful and OEM partitions are
+// mountable with "/dev/ubiX_0". The input is a partition device such as
+// /dev/sda3. Return empty string on error.
+std::string MakePartitionNameForMount(const std::string& part_name);
+
 // Returns the sysfs block device for a root block device. For
 // example, SysfsBlockDevice("/dev/sda") returns
 // "/sys/block/sda". Returns an empty string if the input device is
diff --git a/utils_unittest.cc b/utils_unittest.cc
index 3010738..cd91279 100644
--- a/utils_unittest.cc
+++ b/utils_unittest.cc
@@ -66,14 +66,14 @@
             utils::KernelDeviceOfBootDevice("/dev/mmcblk0p3"));
   EXPECT_EQ("", utils::KernelDeviceOfBootDevice("/dev/mmcblk0p4"));
 
-  EXPECT_EQ("/dev/ubi2", utils::KernelDeviceOfBootDevice("/dev/ubi3"));
+  EXPECT_EQ("/dev/mtd2", utils::KernelDeviceOfBootDevice("/dev/ubi3"));
   EXPECT_EQ("", utils::KernelDeviceOfBootDevice("/dev/ubi4"));
 
-  EXPECT_EQ("/dev/mtdblock2",
+  EXPECT_EQ("/dev/mtd2",
             utils::KernelDeviceOfBootDevice("/dev/ubiblock3_0"));
-  EXPECT_EQ("/dev/mtdblock4",
+  EXPECT_EQ("/dev/mtd4",
             utils::KernelDeviceOfBootDevice("/dev/ubiblock5_0"));
-  EXPECT_EQ("/dev/mtdblock6",
+  EXPECT_EQ("/dev/mtd6",
             utils::KernelDeviceOfBootDevice("/dev/ubiblock7_0"));
   EXPECT_EQ("", utils::KernelDeviceOfBootDevice("/dev/ubiblock4_0"));
 }
@@ -218,7 +218,32 @@
   EXPECT_EQ("/dev/mmcblk0p2", utils::MakePartitionName("/dev/mmcblk0", 2));
   EXPECT_EQ("/dev/loop8", utils::MakePartitionName("/dev/loop", 8));
   EXPECT_EQ("/dev/loop12p2", utils::MakePartitionName("/dev/loop12", 2));
-  EXPECT_EQ("/dev/ubiblock3_0", utils::MakePartitionName("/dev/ubiblock", 3));
+  EXPECT_EQ("/dev/ubi5_0", utils::MakePartitionName("/dev/ubiblock", 5));
+  EXPECT_EQ("/dev/mtd4", utils::MakePartitionName("/dev/ubiblock", 4));
+  EXPECT_EQ("/dev/ubi3_0", utils::MakePartitionName("/dev/ubiblock", 3));
+  EXPECT_EQ("/dev/mtd2", utils::MakePartitionName("/dev/ubiblock", 2));
+  EXPECT_EQ("/dev/ubi1_0", utils::MakePartitionName("/dev/ubiblock", 1));
+}
+
+TEST(UtilsTest, MakePartitionNameForMountTest) {
+  EXPECT_EQ("/dev/sda4", utils::MakePartitionNameForMount("/dev/sda4"));
+  EXPECT_EQ("/dev/sda123", utils::MakePartitionNameForMount("/dev/sda123"));
+  EXPECT_EQ("/dev/mmcblk2", utils::MakePartitionNameForMount("/dev/mmcblk2"));
+  EXPECT_EQ("/dev/mmcblk0p2",
+            utils::MakePartitionNameForMount("/dev/mmcblk0p2"));
+  EXPECT_EQ("/dev/loop0", utils::MakePartitionNameForMount("/dev/loop0"));
+  EXPECT_EQ("/dev/loop8", utils::MakePartitionNameForMount("/dev/loop8"));
+  EXPECT_EQ("/dev/loop12p2",
+            utils::MakePartitionNameForMount("/dev/loop12p2"));
+  EXPECT_EQ("/dev/ubiblock5_0",
+            utils::MakePartitionNameForMount("/dev/ubiblock5_0"));
+  EXPECT_EQ("/dev/mtd4",
+            utils::MakePartitionNameForMount("/dev/ubi4_0"));
+  EXPECT_EQ("/dev/ubiblock3_0",
+            utils::MakePartitionNameForMount("/dev/ubiblock3"));
+  EXPECT_EQ("/dev/mtd2", utils::MakePartitionNameForMount("/dev/ubi2"));
+  EXPECT_EQ("/dev/ubi1_0",
+            utils::MakePartitionNameForMount("/dev/ubiblock1"));
 }
 
 namespace {
@@ -372,11 +397,11 @@
 
   boot_dev = "/dev/ubiblock3_0";
   EXPECT_TRUE(utils::GetInstallDev(boot_dev, &install_dev));
-  EXPECT_EQ(install_dev, "/dev/ubiblock5_0");
+  EXPECT_EQ(install_dev, "/dev/ubi5_0");
 
   boot_dev = "/dev/ubiblock5_0";
   EXPECT_TRUE(utils::GetInstallDev(boot_dev, &install_dev));
-  EXPECT_EQ(install_dev, "/dev/ubiblock3_0");
+  EXPECT_EQ(install_dev, "/dev/ubi3_0");
 
   boot_dev = "/dev/ubiblock12_0";
   EXPECT_FALSE(utils::GetInstallDev(boot_dev, &install_dev));