diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
index 2fe04bd..e4c0934 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -64,8 +64,8 @@
 
     private AtomicBoolean mManualSwitchSlotRequired = new AtomicBoolean(true);
 
-    /** Validate state only once when app binds to UpdateEngine. */
-    private AtomicBoolean mStateValidityEnsured = new AtomicBoolean(false);
+    /** Synchronize state with engine status only once when app binds to UpdateEngine. */
+    private AtomicBoolean mStateSynchronized = new AtomicBoolean(false);
 
     @GuardedBy("mLock")
     private UpdateData mLastUpdateData = null;
@@ -90,10 +90,12 @@
     }
 
     /**
-     * Binds to {@link UpdateEngine}.
+     * Binds to {@link UpdateEngine}. Invokes onStateChangeCallback if present.
      */
     public void bind() {
-        mStateValidityEnsured.set(false);
+        getOnStateChangeCallback().ifPresent(callback -> callback.accept(mUpdaterState.get()));
+
+        mStateSynchronized.set(false);
         this.mUpdateEngine.bind(mUpdateEngineCallback);
     }
 
@@ -104,11 +106,8 @@
         this.mUpdateEngine.unbind();
     }
 
-    /**
-     * @return a number from {@code 0.0} to {@code 1.0}.
-     */
-    public float getProgress() {
-        return (float) this.mProgress.get();
+    public int getUpdaterState() {
+        return mUpdaterState.get();
     }
 
     /**
@@ -202,18 +201,37 @@
      * Updates {@link this.mState} and if state is changed,
      * it also notifies {@link this.mOnStateChangeCallback}.
      */
-    private void setUpdaterState(int updaterState) {
+    private void setUpdaterState(int newUpdaterState)
+            throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "setUpdaterState invoked newState=" + newUpdaterState);
         int previousState = mUpdaterState.get();
+        mUpdaterState.set(newUpdaterState);
+        if (previousState != newUpdaterState) {
+            getOnStateChangeCallback().ifPresent(callback -> callback.accept(newUpdaterState));
+        }
+    }
+
+    /**
+     * Same as {@link this.setUpdaterState}. Logs the error if new state
+     * cannot be set.
+     */
+    private void setUpdaterStateSilent(int newUpdaterState) {
         try {
-            mUpdaterState.set(updaterState);
+            setUpdaterState(newUpdaterState);
         } catch (UpdaterState.InvalidTransitionException e) {
-            // Note: invalid state transitions should be handled properly,
-            //       but to make sample app simple, we just throw runtime exception.
-            throw new RuntimeException("Can't set state " + updaterState, e);
+            // Most likely UpdateEngine status and UpdaterSample state got de-synchronized.
+            // To make sample app simple, we don't handle it properly.
+            Log.e(TAG, "Failed to set updater state", e);
         }
-        if (previousState != updaterState) {
-            getOnStateChangeCallback().ifPresent(callback -> callback.accept(updaterState));
-        }
+    }
+
+    /**
+     * Creates new UpdaterState, assigns it to {@link this.mUpdaterState},
+     * and notifies callbacks.
+     */
+    private void initializeUpdateState(int state) {
+        this.mUpdaterState = new UpdaterState(state);
+        getOnStateChangeCallback().ifPresent(callback -> callback.accept(state));
     }
 
     /**
@@ -224,13 +242,10 @@
      * update engine would throw an error when the method is called, and the only way to
      * handle it is to catch the exception.</p>
      */
-    public void cancelRunningUpdate() {
-        try {
-            mUpdateEngine.cancel();
-            setUpdaterState(UpdaterState.IDLE);
-        } catch (Exception e) {
-            Log.w(TAG, "UpdateEngine failed to stop the ongoing update", e);
-        }
+    public synchronized void cancelRunningUpdate() throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "cancelRunningUpdate invoked");
+        setUpdaterState(UpdaterState.IDLE);
+        mUpdateEngine.cancel();
     }
 
     /**
@@ -240,13 +255,10 @@
      * update engine would throw an error when the method is called, and the only way to
      * handle it is to catch the exception.</p>
      */
-    public void resetUpdate() {
-        try {
-            mUpdateEngine.resetStatus();
-            setUpdaterState(UpdaterState.IDLE);
-        } catch (Exception e) {
-            Log.w(TAG, "UpdateEngine failed to reset the update", e);
-        }
+    public synchronized void resetUpdate() throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "resetUpdate invoked");
+        setUpdaterState(UpdaterState.IDLE);
+        mUpdateEngine.resetStatus();
     }
 
     /**
@@ -255,7 +267,8 @@
      * <p>UpdateEngine works asynchronously. This method doesn't wait until
      * end of the update.</p>
      */
-    public void applyUpdate(Context context, UpdateConfig config) {
+    public synchronized void applyUpdate(Context context, UpdateConfig config)
+            throws UpdaterState.InvalidTransitionException {
         mEngineErrorCode.set(UpdateEngineErrorCodes.UNKNOWN);
         setUpdaterState(UpdaterState.RUNNING);
 
@@ -277,7 +290,8 @@
         }
     }
 
-    private void applyAbNonStreamingUpdate(UpdateConfig config) {
+    private void applyAbNonStreamingUpdate(UpdateConfig config)
+            throws UpdaterState.InvalidTransitionException {
         UpdateData.Builder builder = UpdateData.builder()
                 .setExtraProperties(prepareExtraProperties(config));
 
@@ -306,7 +320,7 @@
                 updateEngineApplyPayload(builder.build());
             } else {
                 Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
-                setUpdaterState(UpdaterState.ERROR);
+                setUpdaterStateSilent(UpdaterState.ERROR);
             }
         });
     }
@@ -333,6 +347,8 @@
      * with the given id.</p>
      */
     private void updateEngineApplyPayload(UpdateData update) {
+        Log.d(TAG, "updateEngineApplyPayload invoked with url " + update.mPayload.getUrl());
+
         synchronized (mLock) {
             mLastUpdateData = update;
         }
@@ -348,11 +364,15 @@
                     properties.toArray(new String[0]));
         } catch (Exception e) {
             Log.e(TAG, "UpdateEngine failed to apply the update", e);
-            setUpdaterState(UpdaterState.ERROR);
+            setUpdaterStateSilent(UpdaterState.ERROR);
         }
     }
 
+    /**
+     * Re-applies {@link this.mLastUpdateData} to update_engine.
+     */
     private void updateEngineReApplyPayload() {
+        Log.d(TAG, "updateEngineReApplyPayload invoked");
         UpdateData lastUpdate;
         synchronized (mLock) {
             // mLastPayloadSpec might be empty in some cases.
@@ -377,8 +397,14 @@
      * {@link UpdateEngine#applyPayload} might take several seconds to finish, and it will
      * invoke callbacks {@link this#onStatusUpdate} and {@link this#onPayloadApplicationComplete)}.
      */
-    public void setSwitchSlotOnReboot() {
+    public synchronized void setSwitchSlotOnReboot() {
         Log.d(TAG, "setSwitchSlotOnReboot invoked");
+
+        // When mManualSwitchSlotRequired set false, next time
+        // onApplicationPayloadComplete is called,
+        // it will set updater state to REBOOT_REQUIRED.
+        mManualSwitchSlotRequired.set(false);
+
         UpdateData.Builder builder;
         synchronized (mLock) {
             // To make sample app simple, we don't handle it.
@@ -396,65 +422,62 @@
     }
 
     /**
-     * Verifies if mUpdaterState matches mUpdateEngineStatus.
-     * If they don't match, runs applyPayload to trigger onPayloadApplicationComplete
-     * callback, which updates mUpdaterState.
+     * Synchronize UpdaterState with UpdateEngine status.
+     * Apply necessary UpdateEngine operation if status are out of sync.
+     *
+     * It's expected to be called once when sample app binds itself to UpdateEngine.
      */
-    private void ensureCorrectUpdaterState() {
-        // When mUpdaterState is one of IDLE, PAUSED, ERROR, SLOT_SWITCH_REQUIRED
-        //    then mUpdateEngineStatus must be IDLE.
-        // When mUpdaterState is RUNNING,
-        //    then mUpdateEngineStatus must not be IDLE or UPDATED_NEED_REBOOT.
-        // When mUpdaterState is REBOOT_REQUIRED,
-        //    then mUpdateEngineStatus must be UPDATED_NEED_REBOOT.
+    private void synchronizeUpdaterStateWithUpdateEngineStatus() {
+        Log.d(TAG, "synchronizeUpdaterStateWithUpdateEngineStatus is invoked.");
+
         int state = mUpdaterState.get();
-        int updateEngineStatus = mUpdateEngineStatus.get();
-        if (state == UpdaterState.IDLE
-                || state == UpdaterState.ERROR
-                || state == UpdaterState.PAUSED
-                || state == UpdaterState.SLOT_SWITCH_REQUIRED) {
-            ensureUpdateEngineStatusIdle(state, updateEngineStatus);
-        } else if (state == UpdaterState.RUNNING) {
-            ensureUpdateEngineStatusRunning(state, updateEngineStatus);
-        } else if (state == UpdaterState.REBOOT_REQUIRED) {
-            ensureUpdateEngineStatusReboot(state, updateEngineStatus);
-        }
-    }
+        int engineStatus = mUpdateEngineStatus.get();
 
-    private void ensureUpdateEngineStatusIdle(int state, int updateEngineStatus) {
-        if (updateEngineStatus == UpdateEngine.UpdateStatusConstants.IDLE) {
+        if (engineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT) {
+            // If update has been installed before running the sample app,
+            // set state to REBOOT_REQUIRED.
+            initializeUpdateState(UpdaterState.REBOOT_REQUIRED);
             return;
         }
-        // It might happen when update is started not from the sample app.
-        // To make the sample app simple, we won't handle this case.
-        throw new RuntimeException("When mUpdaterState is " + state
-                + " mUpdateEngineStatus expected to be "
-                + UpdateEngine.UpdateStatusConstants.IDLE
-                + ", but it is " + updateEngineStatus);
-    }
 
-    private void ensureUpdateEngineStatusRunning(int state, int updateEngineStatus) {
-        if (updateEngineStatus != UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT
-                && updateEngineStatus != UpdateEngine.UpdateStatusConstants.IDLE) {
-            return;
+        switch (state) {
+            case UpdaterState.IDLE:
+            case UpdaterState.ERROR:
+            case UpdaterState.PAUSED:
+            case UpdaterState.SLOT_SWITCH_REQUIRED:
+                // It might happen when update is started not from the sample app.
+                // To make the sample app simple, we won't handle this case.
+                Preconditions.checkState(
+                        engineStatus == UpdateEngine.UpdateStatusConstants.IDLE,
+                        "When mUpdaterState is %s, mUpdateEngineStatus "
+                                + "must be 0/IDLE, but it is %s",
+                        state,
+                        engineStatus);
+                break;
+            case UpdaterState.RUNNING:
+                if (engineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT
+                        || engineStatus == UpdateEngine.UpdateStatusConstants.IDLE) {
+                    Log.i(TAG, "ensureUpdateEngineStatusIsRunning - re-applying last payload");
+                    // Re-apply latest update. It makes update_engine to invoke
+                    // onPayloadApplicationComplete callback. The callback notifies
+                    // if update was successful or not.
+                    updateEngineReApplyPayload();
+                }
+                break;
+            case UpdaterState.REBOOT_REQUIRED:
+                // This might happen when update is installed by other means,
+                // and sample app is not aware of it.
+                // To make the sample app simple, we won't handle this case.
+                Preconditions.checkState(
+                        engineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT,
+                        "When mUpdaterState is %s, mUpdateEngineStatus "
+                                + "must be 6/UPDATED_NEED_REBOOT, but it is %s",
+                        state,
+                        engineStatus);
+                break;
+            default:
+                throw new IllegalStateException("This block should not be reached.");
         }
-        // Re-apply latest update. It makes update_engine to invoke
-        // onPayloadApplicationComplete callback. The callback notifies
-        // if update was successful or not.
-        updateEngineReApplyPayload();
-    }
-
-    private void ensureUpdateEngineStatusReboot(int state, int updateEngineStatus) {
-        if (updateEngineStatus == UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT) {
-            return;
-        }
-        // This might happen when update is installed by other means,
-        // and sample app is not aware of it. To make the sample app simple,
-        // we won't handle this case.
-        throw new RuntimeException("When mUpdaterState is " + state
-                + " mUpdateEngineStatus expected to be "
-                + UpdateEngine.UpdateStatusConstants.UPDATED_NEED_REBOOT
-                + ", but it is " + updateEngineStatus);
     }
 
     /**
@@ -468,13 +491,19 @@
      * @param progress a number from 0.0 to 1.0.
      */
     private void onStatusUpdate(int status, float progress) {
+        Log.d(TAG, String.format(
+                        "onStatusUpdate invoked, status=%s, progress=%.2f",
+                        status,
+                        progress));
+
         int previousStatus = mUpdateEngineStatus.get();
         mUpdateEngineStatus.set(status);
         mProgress.set(progress);
 
-        if (!mStateValidityEnsured.getAndSet(true)) {
-            // We ensure correct state once only when sample app is bound to UpdateEngine.
-            ensureCorrectUpdaterState();
+        if (!mStateSynchronized.getAndSet(true)) {
+            // We synchronize state with engine status once
+            // only when sample app is bound to UpdateEngine.
+            synchronizeUpdaterStateWithUpdateEngineStatus();
         }
 
         getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(progress));
@@ -489,11 +518,11 @@
         mEngineErrorCode.set(errorCode);
         if (errorCode == UpdateEngine.ErrorCodeConstants.SUCCESS
                 || errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
-            setUpdaterState(isManualSwitchSlotRequired()
+            setUpdaterStateSilent(isManualSwitchSlotRequired()
                     ? UpdaterState.SLOT_SWITCH_REQUIRED
                     : UpdaterState.REBOOT_REQUIRED);
         } else if (errorCode != UpdateEngineErrorCodes.USER_CANCELLED) {
-            setUpdaterState(UpdaterState.ERROR);
+            setUpdaterStateSilent(UpdaterState.ERROR);
         }
 
         getOnEngineCompleteCallback()
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
index 36a9098..573d336 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
@@ -51,12 +51,15 @@
      * are allowed to transition to from key.
      */
     private static final ImmutableMap<Integer, ImmutableSet<Integer>> TRANSITIONS =
-            ImmutableMap.of(
-                    IDLE, ImmutableSet.of(RUNNING),
-                    RUNNING, ImmutableSet.of(ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED),
-                    PAUSED, ImmutableSet.of(RUNNING),
-                    SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR)
-            );
+            ImmutableMap.<Integer, ImmutableSet<Integer>>builder()
+                    .put(IDLE, ImmutableSet.of(ERROR, RUNNING))
+                    .put(RUNNING, ImmutableSet.of(
+                            ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED))
+                    .put(PAUSED, ImmutableSet.of(ERROR, RUNNING, IDLE))
+                    .put(SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR, IDLE))
+                    .put(ERROR, ImmutableSet.of(IDLE))
+                    .put(REBOOT_REQUIRED, ImmutableSet.of(IDLE))
+                    .build();
 
     private AtomicInteger mState;
 
diff --git a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
index 1de72c2..6c71cb6 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -88,7 +88,7 @@
 
         this.mTextViewConfigsDirHint.setText(UpdateConfigs.getConfigsRoot(this));
 
-        uiReset();
+        uiResetWidgets();
         loadUpdateConfigs();
 
         this.mUpdateManager.setOnStateChangeCallback(this::onUpdaterStateChange);
@@ -108,7 +108,6 @@
     @Override
     protected void onResume() {
         super.onResume();
-        // TODO(zhomart) load saved states
         // Binding to UpdateEngine invokes onStatusUpdate callback,
         // persisted updater state has to be loaded and prepared beforehand.
         this.mUpdateManager.bind();
@@ -117,7 +116,6 @@
     @Override
     protected void onPause() {
         this.mUpdateManager.unbind();
-        // TODO(zhomart) save state
         super.onPause();
     }
 
@@ -149,14 +147,22 @@
                 .setMessage("Do you really want to apply this update?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    uiSetUpdating();
+                    uiResetWidgets();
                     uiResetEngineText();
-                    mUpdateManager.applyUpdate(this, getSelectedConfig());
+                    applyUpdate(getSelectedConfig());
                 })
                 .setNegativeButton(android.R.string.cancel, null)
                 .show();
     }
 
+    private void applyUpdate(UpdateConfig config) {
+        try {
+            mUpdateManager.applyUpdate(this, config);
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to apply update " + config.getName(), e);
+        }
+    }
+
     /**
      * stop button clicked
      */
@@ -166,11 +172,19 @@
                 .setMessage("Do you really want to cancel running update?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    mUpdateManager.cancelRunningUpdate();
+                    cancelRunningUpdate();
                 })
                 .setNegativeButton(android.R.string.cancel, null).show();
     }
 
+    private void cancelRunningUpdate() {
+        try {
+            mUpdateManager.cancelRunningUpdate();
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to cancel running update", e);
+        }
+    }
+
     /**
      * reset button clicked
      */
@@ -181,11 +195,19 @@
                         + " and restore old version?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    mUpdateManager.resetUpdate();
+                    resetUpdate();
                 })
                 .setNegativeButton(android.R.string.cancel, null).show();
     }
 
+    private void resetUpdate() {
+        try {
+            mUpdateManager.resetUpdate();
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to reset update", e);
+        }
+    }
+
     /**
      * switch slot button clicked
      */
@@ -199,9 +221,25 @@
      * values from {@link UpdaterState}.
      */
     private void onUpdaterStateChange(int state) {
-        Log.i(TAG, "onUpdaterStateChange invoked state=" + state);
+        Log.i(TAG, "UpdaterStateChange state="
+                + UpdaterState.getStateText(state)
+                + "/" + state);
         runOnUiThread(() -> {
             setUiUpdaterState(state);
+
+            if (state == UpdaterState.IDLE) {
+                uiStateIdle();
+            } else if (state == UpdaterState.RUNNING) {
+                uiStateRunning();
+            } else if (state == UpdaterState.PAUSED) {
+                uiStatePaused();
+            } else if (state == UpdaterState.ERROR) {
+                uiStateError();
+            } else if (state == UpdaterState.SLOT_SWITCH_REQUIRED) {
+                uiStateSlotSwitchRequired();
+            } else if (state == UpdaterState.REBOOT_REQUIRED) {
+                uiStateRebootRequired();
+            }
         });
     }
 
@@ -210,17 +248,10 @@
      * be one of the values from {@link UpdateEngine.UpdateStatusConstants}.
      */
     private void onEngineStatusUpdate(int status) {
+        Log.i(TAG, "StatusUpdate - status="
+                + UpdateEngineStatuses.getStatusText(status)
+                + "/" + status);
         runOnUiThread(() -> {
-            Log.e(TAG, "StatusUpdate - status="
-                    + UpdateEngineStatuses.getStatusText(status)
-                    + "/" + status);
-            if (status == UpdateEngine.UpdateStatusConstants.IDLE) {
-                Log.d(TAG, "status changed, resetting ui");
-                uiReset();
-            } else {
-                Log.d(TAG, "status changed, setting ui to updating mode");
-                uiSetUpdating();
-            }
             setUiEngineStatus(status);
         });
     }
@@ -234,19 +265,12 @@
         final String completionState = UpdateEngineErrorCodes.isUpdateSucceeded(errorCode)
                 ? "SUCCESS"
                 : "FAILURE";
+        Log.i(TAG,
+                "PayloadApplicationCompleted - errorCode="
+                        + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
+                        + " " + completionState);
         runOnUiThread(() -> {
-            Log.i(TAG,
-                    "Completed - errorCode="
-                            + UpdateEngineErrorCodes.getCodeName(errorCode) + "/" + errorCode
-                            + " " + completionState);
             setUiEngineErrorCode(errorCode);
-            if (errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
-                // if update was successfully applied.
-                if (mUpdateManager.isManualSwitchSlotRequired()) {
-                    // Show "Switch Slot" button.
-                    uiShowSwitchSlotInfo();
-                }
-            }
         });
     }
 
@@ -258,45 +282,66 @@
     }
 
     /** resets ui */
-    private void uiReset() {
+    private void uiResetWidgets() {
         mTextViewBuild.setText(Build.DISPLAY);
-        mSpinnerConfigs.setEnabled(true);
-        mButtonReload.setEnabled(true);
-        mButtonApplyConfig.setEnabled(true);
+        mSpinnerConfigs.setEnabled(false);
+        mButtonReload.setEnabled(false);
+        mButtonApplyConfig.setEnabled(false);
         mButtonStop.setEnabled(false);
         mButtonReset.setEnabled(false);
-        mProgressBar.setProgress(0);
         mProgressBar.setEnabled(false);
         mProgressBar.setVisibility(ProgressBar.INVISIBLE);
-        uiHideSwitchSlotInfo();
+        mButtonSwitchSlot.setEnabled(false);
+        mTextViewUpdateInfo.setTextColor(Color.parseColor("#aaaaaa"));
     }
 
     private void uiResetEngineText() {
         mTextViewEngineStatus.setText(R.string.unknown);
         mTextViewEngineErrorCode.setText(R.string.unknown);
-        // Note: Do not reset mTextViewUpdaterState; UpdateManager notifies properly.
+        // Note: Do not reset mTextViewUpdaterState; UpdateManager notifies updater state properly.
     }
 
-    /** sets ui updating mode */
-    private void uiSetUpdating() {
-        mTextViewBuild.setText(Build.DISPLAY);
-        mSpinnerConfigs.setEnabled(false);
-        mButtonReload.setEnabled(false);
-        mButtonApplyConfig.setEnabled(false);
-        mButtonStop.setEnabled(true);
+    private void uiStateIdle() {
+        uiResetWidgets();
+        mSpinnerConfigs.setEnabled(true);
+        mButtonReload.setEnabled(true);
+        mButtonApplyConfig.setEnabled(true);
+        mProgressBar.setProgress(0);
+    }
+
+    private void uiStateRunning() {
+        uiResetWidgets();
         mProgressBar.setEnabled(true);
+        mProgressBar.setVisibility(ProgressBar.VISIBLE);
+        mButtonStop.setEnabled(true);
+    }
+
+    private void uiStatePaused() {
+        uiResetWidgets();
         mButtonReset.setEnabled(true);
+        mProgressBar.setEnabled(true);
         mProgressBar.setVisibility(ProgressBar.VISIBLE);
     }
 
-    private void uiShowSwitchSlotInfo() {
+    private void uiStateSlotSwitchRequired() {
+        uiResetWidgets();
+        mButtonReset.setEnabled(true);
+        mProgressBar.setEnabled(true);
+        mProgressBar.setVisibility(ProgressBar.VISIBLE);
         mButtonSwitchSlot.setEnabled(true);
         mTextViewUpdateInfo.setTextColor(Color.parseColor("#777777"));
     }
 
-    private void uiHideSwitchSlotInfo() {
-        mTextViewUpdateInfo.setTextColor(Color.parseColor("#AAAAAA"));
-        mButtonSwitchSlot.setEnabled(false);
+    private void uiStateError() {
+        uiResetWidgets();
+        mButtonReset.setEnabled(true);
+        mProgressBar.setEnabled(true);
+        mProgressBar.setVisibility(ProgressBar.VISIBLE);
+    }
+
+    private void uiStateRebootRequired() {
+        uiResetWidgets();
+        mButtonReset.setEnabled(true);
     }
 
     /**
diff --git a/updater_sample/tests/res/raw/update_config_stream_001.json b/updater_sample/tests/res/raw/update_config_001_stream.json
similarity index 100%
rename from updater_sample/tests/res/raw/update_config_stream_001.json
rename to updater_sample/tests/res/raw/update_config_001_stream.json
diff --git a/updater_sample/tests/res/raw/update_config_stream_002.json b/updater_sample/tests/res/raw/update_config_002_stream.json
similarity index 100%
rename from updater_sample/tests/res/raw/update_config_stream_002.json
rename to updater_sample/tests/res/raw/update_config_002_stream.json
diff --git a/updater_sample/tests/res/raw/update_config_003_nonstream.json b/updater_sample/tests/res/raw/update_config_003_nonstream.json
new file mode 100644
index 0000000..4175c35
--- /dev/null
+++ b/updater_sample/tests/res/raw/update_config_003_nonstream.json
@@ -0,0 +1,9 @@
+{
+    "__": "*** Generated using tools/gen_update_config.py ***",
+    "ab_config": {
+        "force_switch_slot": false
+    },
+    "ab_install_type": "NON_STREAMING",
+    "name": "S ota_002_package",
+    "url": "file:///data/sample-ota-packages/ota_003_package.zip"
+}
\ No newline at end of file
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
index 000f566..1cbd860 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateConfigTest.java
@@ -60,7 +60,7 @@
     public void setUp() throws Exception {
         mContext = InstrumentationRegistry.getContext();
         mTargetContext = InstrumentationRegistry.getTargetContext();
-        mJsonStreaming001 = readResource(R.raw.update_config_stream_001);
+        mJsonStreaming001 = readResource(R.raw.update_config_001_stream);
     }
 
     @Test
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
index 0657a5e..e05ad29 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
@@ -18,19 +18,22 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
-import static org.mockito.ArgumentMatchers.anyInt;
-import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.content.Context;
 import android.os.UpdateEngine;
 import android.os.UpdateEngineCallback;
+import android.support.test.InstrumentationRegistry;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
 
+import com.example.android.systemupdatersample.tests.R;
 import com.example.android.systemupdatersample.util.PayloadSpecs;
+import com.google.common.collect.ImmutableList;
+import com.google.common.io.CharStreams;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -40,7 +43,9 @@
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
-import java.util.function.IntConsumer;
+import java.io.File;
+import java.io.IOException;
+import java.io.InputStreamReader;
 
 /**
  * Tests for {@link UpdateManager}
@@ -56,37 +61,86 @@
     private UpdateEngine mUpdateEngine;
     @Mock
     private PayloadSpecs mPayloadSpecs;
-    private UpdateManager mUpdateManager;
+    private UpdateManager mSubject;
+    private Context mContext;
+    private UpdateConfig mNonStreamingUpdate003;
 
     @Before
-    public void setUp() {
-        mUpdateManager = new UpdateManager(mUpdateEngine, mPayloadSpecs);
+    public void setUp() throws Exception {
+        mContext = InstrumentationRegistry.getContext();
+        mSubject = new UpdateManager(mUpdateEngine, mPayloadSpecs);
+        mNonStreamingUpdate003 =
+                UpdateConfig.fromJson(readResource(R.raw.update_config_003_nonstream));
     }
 
     @Test
-    public void storesProgressThenInvokesCallbacks() {
-        IntConsumer statusUpdateCallback = mock(IntConsumer.class);
-
-        // When UpdateManager is bound to update_engine, it passes
-        // UpdateManager.UpdateEngineCallbackImpl as a callback to update_engine.
+    public void applyUpdate_appliesPayloadToUpdateEngine() throws Exception {
+        PayloadSpec payload = buildMockPayloadSpec();
+        when(mPayloadSpecs.forNonStreaming(any(File.class))).thenReturn(payload);
         when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
+            // When UpdateManager is bound to update_engine, it passes
+            // UpdateEngineCallback as a callback to update_engine.
             UpdateEngineCallback callback = answer.getArgument(0);
-            callback.onStatusUpdate(/*engineStatus*/ 4, /*engineProgress*/ 0.2f);
+            callback.onStatusUpdate(
+                    UpdateEngine.UpdateStatusConstants.IDLE,
+                    /*engineProgress*/ 0.0f);
             return null;
         });
 
-        mUpdateManager.setOnEngineStatusUpdateCallback(statusUpdateCallback);
+        mSubject.bind();
+        mSubject.applyUpdate(null, mNonStreamingUpdate003);
 
-        // Making sure that manager.getProgress() returns correct progress
-        // in "onEngineStatusUpdate" callback.
-        doAnswer(answer -> {
-            assertEquals(0.2f, mUpdateManager.getProgress(), 1E-5);
+        verify(mUpdateEngine).applyPayload(
+                "file://blah",
+                120,
+                340,
+                new String[] {
+                        "SWITCH_SLOT_ON_REBOOT=0" // ab_config.force_switch_slot = false
+                });
+    }
+
+    @Test
+    public void stateIsRunningAndEngineStatusIsIdle_reApplyLastUpdate() throws Exception {
+        PayloadSpec payload = buildMockPayloadSpec();
+        when(mPayloadSpecs.forNonStreaming(any(File.class))).thenReturn(payload);
+        when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
+            // When UpdateManager is bound to update_engine, it passes
+            // UpdateEngineCallback as a callback to update_engine.
+            UpdateEngineCallback callback = answer.getArgument(0);
+            callback.onStatusUpdate(
+                    UpdateEngine.UpdateStatusConstants.IDLE,
+                    /*engineProgress*/ 0.0f);
             return null;
-        }).when(statusUpdateCallback).accept(anyInt());
+        });
 
-        mUpdateManager.bind();
+        mSubject.bind();
+        mSubject.applyUpdate(null, mNonStreamingUpdate003);
+        mSubject.unbind();
+        mSubject.bind(); // re-bind - now it should re-apply last update
 
-        verify(statusUpdateCallback, times(1)).accept(4);
+        assertEquals(mSubject.getUpdaterState(), UpdaterState.RUNNING);
+        // it should be called 2 times
+        verify(mUpdateEngine, times(2)).applyPayload(
+                "file://blah",
+                120,
+                340,
+                new String[] {
+                        "SWITCH_SLOT_ON_REBOOT=0" // ab_config.force_switch_slot = false
+                });
+    }
+
+    private PayloadSpec buildMockPayloadSpec() {
+        PayloadSpec payload = mock(PayloadSpec.class);
+        when(payload.getUrl()).thenReturn("file://blah");
+        when(payload.getOffset()).thenReturn(120L);
+        when(payload.getSize()).thenReturn(340L);
+        when(payload.getProperties()).thenReturn(ImmutableList.of());
+        return payload;
+    }
+
+    private String readResource(int id) throws IOException {
+        return CharStreams.toString(new InputStreamReader(
+                mContext.getResources().openRawResource(id)));
     }
 
 }