diff --git a/frontend/afe/moblab_rpc_interface.py b/frontend/afe/moblab_rpc_interface.py
index 50972a0..cefb488 100644
--- a/frontend/afe/moblab_rpc_interface.py
+++ b/frontend/afe/moblab_rpc_interface.py
@@ -518,14 +518,10 @@
     # then check if a reboot is needed
     try:
         subprocess.check_call(['sudo', _UPDATE_ENGINE_CLIENT, '--update'])
-        try:
-            # --is_reboot_needed returns 2 if a reboot is required, which
-            # technically is an error
-            subprocess.check_call(
-                ['sudo', _UPDATE_ENGINE_CLIENT, '--is_reboot_needed'])
-        except subprocess.CalledProcessError as e:
-            if e.returncode == 2:
-                subprocess.call(['sudo', _UPDATE_ENGINE_CLIENT, '--reboot'])
+        # --is_reboot_needed returns 0 if a reboot is required
+        subprocess.check_call(
+            ['sudo', _UPDATE_ENGINE_CLIENT, '--is_reboot_needed'])
+        subprocess.call(['sudo', _UPDATE_ENGINE_CLIENT, '--reboot'])
 
     except subprocess.CalledProcessError as e:
         pass
diff --git a/frontend/afe/moblab_rpc_interface_unittest.py b/frontend/afe/moblab_rpc_interface_unittest.py
index 484a666..87dc346 100644
--- a/frontend/afe/moblab_rpc_interface_unittest.py
+++ b/frontend/afe/moblab_rpc_interface_unittest.py
@@ -496,9 +496,8 @@
         self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'check_call')
         moblab_rpc_interface.subprocess.check_call(['sudo',
                 update_engine_client, '--update'])
-        error = moblab_rpc_interface.subprocess.CalledProcessError(2, '')
         moblab_rpc_interface.subprocess.check_call(['sudo',
-                update_engine_client, '--is_reboot_needed']).AndRaise(error)
+                update_engine_client, '--is_reboot_needed'])
 
         self.mox.StubOutWithMock(moblab_rpc_interface.subprocess, 'call')
         moblab_rpc_interface.subprocess.call(['sudo', update_engine_client,
diff --git a/frontend/client/src/autotest/moblab/wizard/ConfigWizard.java b/frontend/client/src/autotest/moblab/wizard/ConfigWizard.java
index ba6d619..77d0484 100644
--- a/frontend/client/src/autotest/moblab/wizard/ConfigWizard.java
+++ b/frontend/client/src/autotest/moblab/wizard/ConfigWizard.java
@@ -132,7 +132,8 @@
           @Override
           public void onClick(ClickEvent event) {
             String windowText = "If an update is available, the device will be "
-              + "rebooted and all running jobs will be halted. Proceed?";
+              + "rebooted and all running jobs will be halted. "
+              + "This may take 5-10 minutes based on network speed. Proceed?";
             if (Window.confirm(windowText)) {
               MoblabRpcHelper.updateMoblab(new JsonRpcCallback() {
                 @Override