diff --git a/install/wipe_device.cpp b/install/wipe_device.cpp
index 0a525fa..2656580 100644
--- a/install/wipe_device.cpp
+++ b/install/wipe_device.cpp
@@ -182,13 +182,17 @@
     LOG(ERROR) << "Failed to open wipe package";
     return false;
   }
+  return WipeAbDevice(device, wipe_package.get());
+}
 
-  if (!CheckWipePackage(wipe_package.get(), ui)) {
+bool WipeAbDevice(Device* device, Package* wipe_package) {
+  auto ui = device->GetUI();
+  if (!CheckWipePackage(wipe_package, ui)) {
     LOG(ERROR) << "Failed to verify wipe package";
     return false;
   }
 
-  auto partition_list = GetWipePartitionList(wipe_package.get());
+  auto partition_list = GetWipePartitionList(wipe_package);
   if (partition_list.empty()) {
     LOG(ERROR) << "Empty wipe ab partition list";
     return false;