3 files changed, 57 insertions, 13 deletions

diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
index e4c09346b..a9783e70a 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -198,6 +198,24 @@ public class UpdateManager {
     }
 
     /**
+     * Suspend running update.
+     */
+    public synchronized void suspend() throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "suspend invoked");
+        setUpdaterState(UpdaterState.PAUSED);
+        mUpdateEngine.cancel();
+    }
+
+    /**
+     * Resume suspended update.
+     */
+    public synchronized void resume() throws UpdaterState.InvalidTransitionException {
+        Log.d(TAG, "resume invoked");
+        setUpdaterState(UpdaterState.RUNNING);
+        updateEngineReApplyPayload();
+    }
+
+    /**
      * Updates {@link this.mState} and if state is changed,
      * it also notifies {@link this.mOnStateChangeCallback}.
      */
@@ -237,10 +255,6 @@ public class UpdateManager {
     /**
      * Requests update engine to stop any ongoing update. If an update has been applied,
      * leave it as is.
-     *
-     * <p>Sometimes it's possible that the
-     * 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 synchronized void cancelRunningUpdate() throws UpdaterState.InvalidTransitionException {
         Log.d(TAG, "cancelRunningUpdate invoked");
@@ -250,10 +264,6 @@ public class UpdateManager {
 
     /**
      * Resets update engine to IDLE state. If an update has been applied it reverts it.
-     *
-     * <p>Sometimes it's possible that the
-     * 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 synchronized void resetUpdate() throws UpdaterState.InvalidTransitionException {
         Log.d(TAG, "resetUpdate invoked");
@@ -506,7 +516,7 @@ public class UpdateManager {
             synchronizeUpdaterStateWithUpdateEngineStatus();
         }
 
-        getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(progress));
+        getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(mProgress.get()));
 
         if (previousStatus != status) {
             getOnEngineStatusUpdateCallback().ifPresent(callback -> callback.accept(status));
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
index 573d336e9..4eb0b68c7 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdaterState.java
@@ -52,12 +52,12 @@ public class UpdaterState {
      */
     private static final ImmutableMap<Integer, ImmutableSet<Integer>> TRANSITIONS =
             ImmutableMap.<Integer, ImmutableSet<Integer>>builder()
-                    .put(IDLE, ImmutableSet.of(ERROR, RUNNING))
+                    .put(IDLE, ImmutableSet.of(IDLE, ERROR, RUNNING))
+                    .put(ERROR, ImmutableSet.of(IDLE))
                     .put(RUNNING, ImmutableSet.of(
-                            ERROR, PAUSED, REBOOT_REQUIRED, SLOT_SWITCH_REQUIRED))
+                            IDLE, 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(SLOT_SWITCH_REQUIRED, ImmutableSet.of(ERROR, REBOOT_REQUIRED, IDLE))
                     .put(REBOOT_REQUIRED, ImmutableSet.of(IDLE))
                     .build();
 
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 6c71cb6f4..fc9fddd70 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -55,6 +55,8 @@ public class MainActivity extends Activity {
     private Button mButtonApplyConfig;
     private Button mButtonStop;
     private Button mButtonReset;
+    private Button mButtonSuspend;
+    private Button mButtonResume;
     private ProgressBar mProgressBar;
     private TextView mTextViewUpdaterState;
     private TextView mTextViewEngineStatus;
@@ -79,6 +81,8 @@ public class MainActivity extends Activity {
         this.mButtonApplyConfig = findViewById(R.id.buttonApplyConfig);
         this.mButtonStop = findViewById(R.id.buttonStop);
         this.mButtonReset = findViewById(R.id.buttonReset);
+        this.mButtonSuspend = findViewById(R.id.buttonSuspend);
+        this.mButtonResume = findViewById(R.id.buttonResume);
         this.mProgressBar = findViewById(R.id.progressBar);
         this.mTextViewUpdaterState = findViewById(R.id.textViewUpdaterState);
         this.mTextViewEngineStatus = findViewById(R.id.textViewEngineStatus);
@@ -209,9 +213,34 @@ public class MainActivity extends Activity {
     }
 
     /**
+     * suspend button clicked
+     */
+    public void onSuspendClick(View view) {
+        try {
+            mUpdateManager.suspend();
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to suspend running update", e);
+        }
+    }
+
+    /**
+     * resume button clicked
+     */
+    public void onResumeClick(View view) {
+        try {
+            uiResetWidgets();
+            uiResetEngineText();
+            mUpdateManager.resume();
+        } catch (UpdaterState.InvalidTransitionException e) {
+            Log.e(TAG, "Failed to resume running update", e);
+        }
+    }
+
+    /**
      * switch slot button clicked
      */
     public void onSwitchSlotClick(View view) {
+        uiResetWidgets();
         mUpdateManager.setSwitchSlotOnReboot();
     }
 
@@ -289,6 +318,8 @@ public class MainActivity extends Activity {
         mButtonApplyConfig.setEnabled(false);
         mButtonStop.setEnabled(false);
         mButtonReset.setEnabled(false);
+        mButtonSuspend.setEnabled(false);
+        mButtonResume.setEnabled(false);
         mProgressBar.setEnabled(false);
         mProgressBar.setVisibility(ProgressBar.INVISIBLE);
         mButtonSwitchSlot.setEnabled(false);
@@ -303,6 +334,7 @@ public class MainActivity extends Activity {
 
     private void uiStateIdle() {
         uiResetWidgets();
+        mButtonReset.setEnabled(true);
         mSpinnerConfigs.setEnabled(true);
         mButtonReload.setEnabled(true);
         mButtonApplyConfig.setEnabled(true);
@@ -314,6 +346,7 @@ public class MainActivity extends Activity {
         mProgressBar.setEnabled(true);
         mProgressBar.setVisibility(ProgressBar.VISIBLE);
         mButtonStop.setEnabled(true);
+        mButtonSuspend.setEnabled(true);
     }
 
     private void uiStatePaused() {
@@ -321,6 +354,7 @@ public class MainActivity extends Activity {
         mButtonReset.setEnabled(true);
         mProgressBar.setEnabled(true);
         mProgressBar.setVisibility(ProgressBar.VISIBLE);
+        mButtonResume.setEnabled(true);
     }
 
     private void uiStateSlotSwitchRequired() {