9 files changed, 492 insertions, 243 deletions

diff --git a/updater_sample/README.md b/updater_sample/README.md
index c68c07caf..3f211ddba 100644
--- a/updater_sample/README.md
+++ b/updater_sample/README.md
@@ -90,9 +90,9 @@ which HTTP headers are supported.
 - [x] Add demo for passing HTTP headers to `UpdateEngine#applyPayload`
 - [x] [Package compatibility check](https://source.android.com/devices/architecture/vintf/match-rules)
 - [x] Deferred switch slot demo
-- [ ] Add tests for `MainActivity`
+- [ ] Add demo for passing NETWORK_ID to `UpdateEngine#applyPayload`
 - [ ] Verify system partition checksum for package
-- [ ] Add non-A/B updates demo
+- [?] Add non-A/B updates demo
 
 
 ## Running tests
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
index db99f7c74..1e0fadc27 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateConfig.java
@@ -279,4 +279,4 @@ public class UpdateConfig implements Parcelable {
 
     }
 
-}
\ No newline at end of file
+}
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
new file mode 100644
index 000000000..9f0a04e33
--- /dev/null
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -0,0 +1,345 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.android.systemupdatersample;
+
+import android.content.Context;
+import android.os.UpdateEngine;
+import android.os.UpdateEngineCallback;
+import android.util.Log;
+
+import com.example.android.systemupdatersample.services.PrepareStreamingService;
+import com.example.android.systemupdatersample.util.PayloadSpecs;
+import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
+import com.example.android.systemupdatersample.util.UpdateEngineProperties;
+import com.google.common.util.concurrent.AtomicDouble;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.DoubleConsumer;
+import java.util.function.IntConsumer;
+
+/**
+ * Manages the update flow. Asynchronously interacts with the {@link UpdateEngine}.
+ */
+public class UpdateManager {
+
+    private static final String TAG = "UpdateManager";
+
+    /** HTTP Header: User-Agent; it will be sent to the server when streaming the payload. */
+    private static final String HTTP_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+            + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/60.0.3112.113 Safari/537.36";
+
+    private final UpdateEngine mUpdateEngine;
+    private final PayloadSpecs mPayloadSpecs;
+
+    private AtomicInteger mUpdateEngineStatus =
+            new AtomicInteger(UpdateEngine.UpdateStatusConstants.IDLE);
+    private AtomicInteger mEngineErrorCode = new AtomicInteger(UpdateEngineErrorCodes.UNKNOWN);
+    private AtomicDouble mProgress = new AtomicDouble(0);
+
+    private final UpdateManager.UpdateEngineCallbackImpl
+            mUpdateEngineCallback = new UpdateManager.UpdateEngineCallbackImpl();
+
+    private PayloadSpec mLastPayloadSpec;
+    private AtomicBoolean mManualSwitchSlotRequired = new AtomicBoolean(true);
+
+    private IntConsumer mOnEngineStatusUpdateCallback = null;
+    private DoubleConsumer mOnProgressUpdateCallback = null;
+    private IntConsumer mOnEngineCompleteCallback = null;
+
+    private final Object mLock = new Object();
+
+    public UpdateManager(UpdateEngine updateEngine, PayloadSpecs payloadSpecs) {
+        this.mUpdateEngine = updateEngine;
+        this.mPayloadSpecs = payloadSpecs;
+    }
+
+    /**
+     * Binds to {@link UpdateEngine}.
+     */
+    public void bind() {
+        this.mUpdateEngine.bind(mUpdateEngineCallback);
+    }
+
+    /**
+     * Unbinds from {@link UpdateEngine}.
+     */
+    public void unbind() {
+        this.mUpdateEngine.unbind();
+    }
+
+    /**
+     * @return a number from {@code 0.0} to {@code 1.0}.
+     */
+    public float getProgress() {
+        return (float) this.mProgress.get();
+    }
+
+    /**
+     * Returns true if manual switching slot is required. Value depends on
+     * the update config {@code ab_config.force_switch_slot}.
+     */
+    public boolean manualSwitchSlotRequired() {
+        return mManualSwitchSlotRequired.get();
+    }
+
+    /**
+     * Sets update engine status update callback. Value of {@code status} will
+     * be one of the values from {@link UpdateEngine.UpdateStatusConstants}.
+     *
+     * @param onStatusUpdateCallback a callback with parameter {@code status}.
+     */
+    public void setOnEngineStatusUpdateCallback(IntConsumer onStatusUpdateCallback) {
+        synchronized (mLock) {
+            this.mOnEngineStatusUpdateCallback = onStatusUpdateCallback;
+        }
+    }
+
+    private Optional<IntConsumer> getOnEngineStatusUpdateCallback() {
+        synchronized (mLock) {
+            return mOnEngineStatusUpdateCallback == null
+                    ? Optional.empty()
+                    : Optional.of(mOnEngineStatusUpdateCallback);
+        }
+    }
+
+    /**
+     * Sets update engine payload application complete callback. Value of {@code errorCode} will
+     * be one of the values from {@link UpdateEngine.ErrorCodeConstants}.
+     *
+     * @param onComplete a callback with parameter {@code errorCode}.
+     */
+    public void setOnEngineCompleteCallback(IntConsumer onComplete) {
+        synchronized (mLock) {
+            this.mOnEngineCompleteCallback = onComplete;
+        }
+    }
+
+    private Optional<IntConsumer> getOnEngineCompleteCallback() {
+        synchronized (mLock) {
+            return mOnEngineCompleteCallback == null
+                    ? Optional.empty()
+                    : Optional.of(mOnEngineCompleteCallback);
+        }
+    }
+
+    /**
+     * Sets progress update callback. Progress is a number from {@code 0.0} to {@code 1.0}.
+     *
+     * @param onProgressCallback a callback with parameter {@code progress}.
+     */
+    public void setOnProgressUpdateCallback(DoubleConsumer onProgressCallback) {
+        synchronized (mLock) {
+            this.mOnProgressUpdateCallback = onProgressCallback;
+        }
+    }
+
+    private Optional<DoubleConsumer> getOnProgressUpdateCallback() {
+        synchronized (mLock) {
+            return mOnProgressUpdateCallback == null
+                    ? Optional.empty()
+                    : Optional.of(mOnProgressUpdateCallback);
+        }
+    }
+
+    /**
+     * 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 void cancelRunningUpdate() {
+        try {
+            mUpdateEngine.cancel();
+        } catch (Exception e) {
+            Log.w(TAG, "UpdateEngine failed to stop the ongoing update", e);
+        }
+    }
+
+    /**
+     * 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 void resetUpdate() {
+        try {
+            mUpdateEngine.resetStatus();
+        } catch (Exception e) {
+            Log.w(TAG, "UpdateEngine failed to reset the update", e);
+        }
+    }
+
+    /**
+     * Applies the given update.
+     *
+     * <p>UpdateEngine works asynchronously. This method doesn't wait until
+     * end of the update.</p>
+     */
+    public void applyUpdate(Context context, UpdateConfig config) {
+        mEngineErrorCode.set(UpdateEngineErrorCodes.UNKNOWN);
+
+        if (!config.getAbConfig().getForceSwitchSlot()) {
+            mManualSwitchSlotRequired.set(true);
+        } else {
+            mManualSwitchSlotRequired.set(false);
+        }
+
+        if (config.getInstallType() == UpdateConfig.AB_INSTALL_TYPE_NON_STREAMING) {
+            applyAbNonStreamingUpdate(config);
+        } else {
+            applyAbStreamingUpdate(context, config);
+        }
+    }
+
+    private void applyAbNonStreamingUpdate(UpdateConfig config) {
+        List<String> extraProperties = prepareExtraProperties(config);
+
+        PayloadSpec payload;
+        try {
+            payload = mPayloadSpecs.forNonStreaming(config.getUpdatePackageFile());
+        } catch (IOException e) {
+            Log.e(TAG, "Error creating payload spec", e);
+            return;
+        }
+        updateEngineApplyPayload(payload, extraProperties);
+    }
+
+    private void applyAbStreamingUpdate(Context context, UpdateConfig config) {
+        List<String> extraProperties = prepareExtraProperties(config);
+
+        Log.d(TAG, "Starting PrepareStreamingService");
+        PrepareStreamingService.startService(context, config, (code, payloadSpec) -> {
+            if (code == PrepareStreamingService.RESULT_CODE_SUCCESS) {
+                extraProperties.add("USER_AGENT=" + HTTP_USER_AGENT);
+                config.getStreamingMetadata()
+                        .getAuthorization()
+                        .ifPresent(s -> extraProperties.add("AUTHORIZATION=" + s));
+                updateEngineApplyPayload(payloadSpec, extraProperties);
+            } else {
+                Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
+            }
+        });
+    }
+
+    private List<String> prepareExtraProperties(UpdateConfig config) {
+        List<String> extraProperties = new ArrayList<>();
+
+        if (!config.getAbConfig().getForceSwitchSlot()) {
+            // Disable switch slot on reboot, which is enabled by default.
+            // User will enable it manually by clicking "Switch Slot" button on the screen.
+            extraProperties.add(UpdateEngineProperties.PROPERTY_DISABLE_SWITCH_SLOT_ON_REBOOT);
+        }
+        return extraProperties;
+    }
+
+    /**
+     * Applies given payload.
+     *
+     * <p>UpdateEngine works asynchronously. This method doesn't wait until
+     * end of the update.</p>
+     *
+     * <p>It's possible that the update engine throws a generic error, such as upon seeing invalid
+     * payload properties (which come from OTA packages), or failing to set up the network
+     * with the given id.</p>
+     *
+     * @param payloadSpec contains url, offset and size to {@code PAYLOAD_BINARY_FILE_NAME}
+     * @param extraProperties additional properties to pass to {@link UpdateEngine#applyPayload}
+     */
+    private void updateEngineApplyPayload(PayloadSpec payloadSpec, List<String> extraProperties) {
+        mLastPayloadSpec = payloadSpec;
+
+        ArrayList<String> properties = new ArrayList<>(payloadSpec.getProperties());
+        if (extraProperties != null) {
+            properties.addAll(extraProperties);
+        }
+        try {
+            mUpdateEngine.applyPayload(
+                    payloadSpec.getUrl(),
+                    payloadSpec.getOffset(),
+                    payloadSpec.getSize(),
+                    properties.toArray(new String[0]));
+        } catch (Exception e) {
+            Log.e(TAG, "UpdateEngine failed to apply the update", e);
+        }
+    }
+
+    /**
+     * Sets the new slot that has the updated partitions as the active slot,
+     * which device will boot into next time.
+     * This method is only supposed to be called after the payload is applied.
+     *
+     * Invoking {@link UpdateEngine#applyPayload} with the same payload url, offset, size
+     * and payload metadata headers doesn't trigger new update. It can be used to just switch
+     * active A/B slot.
+     *
+     * {@link UpdateEngine#applyPayload} might take several seconds to finish, and it will
+     * invoke callbacks {@link this#onStatusUpdate} and {@link this#onPayloadApplicationComplete)}.
+     */
+    public void setSwitchSlotOnReboot() {
+        Log.d(TAG, "setSwitchSlotOnReboot invoked");
+        List<String> extraProperties = new ArrayList<>();
+        // PROPERTY_SKIP_POST_INSTALL should be passed on to skip post-installation hooks.
+        extraProperties.add(UpdateEngineProperties.PROPERTY_SKIP_POST_INSTALL);
+        // It sets property SWITCH_SLOT_ON_REBOOT=1 by default.
+        // HTTP headers are not required, UpdateEngine is not expected to stream payload.
+        updateEngineApplyPayload(mLastPayloadSpec, extraProperties);
+    }
+
+    private void onStatusUpdate(int status, float progress) {
+        int previousStatus = mUpdateEngineStatus.get();
+        mUpdateEngineStatus.set(status);
+        mProgress.set(progress);
+
+        getOnProgressUpdateCallback().ifPresent(callback -> callback.accept(progress));
+
+        if (previousStatus != status) {
+            getOnEngineStatusUpdateCallback().ifPresent(callback -> callback.accept(status));
+        }
+    }
+
+    private void onPayloadApplicationComplete(int errorCode) {
+        Log.d(TAG, "onPayloadApplicationComplete invoked, errorCode=" + errorCode);
+        mEngineErrorCode.set(errorCode);
+
+        getOnEngineCompleteCallback()
+                .ifPresent(callback -> callback.accept(errorCode));
+    }
+
+    /**
+     * Helper class to delegate {@code update_engine} callbacks to UpdateManager
+     */
+    class UpdateEngineCallbackImpl extends UpdateEngineCallback {
+        @Override
+        public void onStatusUpdate(int status, float percent) {
+            UpdateManager.this.onStatusUpdate(status, percent);
+        }
+
+        @Override
+        public void onPayloadApplicationComplete(int errorCode) {
+            UpdateManager.this.onPayloadApplicationComplete(errorCode);
+        }
+    }
+
+}
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 9bab1319d..9237bc794 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -22,7 +22,6 @@ import android.graphics.Color;
 import android.os.Build;
 import android.os.Bundle;
 import android.os.UpdateEngine;
-import android.os.UpdateEngineCallback;
 import android.util.Log;
 import android.view.View;
 import android.widget.ArrayAdapter;
@@ -32,21 +31,15 @@ import android.widget.Spinner;
 import android.widget.TextView;
 import android.widget.Toast;
 
-import com.example.android.systemupdatersample.PayloadSpec;
 import com.example.android.systemupdatersample.R;
 import com.example.android.systemupdatersample.UpdateConfig;
-import com.example.android.systemupdatersample.services.PrepareStreamingService;
+import com.example.android.systemupdatersample.UpdateManager;
 import com.example.android.systemupdatersample.util.PayloadSpecs;
 import com.example.android.systemupdatersample.util.UpdateConfigs;
 import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
-import com.example.android.systemupdatersample.util.UpdateEngineProperties;
 import com.example.android.systemupdatersample.util.UpdateEngineStatuses;
 
-import java.io.IOException;
-import java.util.ArrayList;
 import java.util.List;
-import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.concurrent.atomic.AtomicInteger;
 
 /**
  * UI for SystemUpdaterSample app.
@@ -55,10 +48,6 @@ public class MainActivity extends Activity {
 
     private static final String TAG = "MainActivity";
 
-    /** HTTP Header: User-Agent; it will be sent to the server when streaming the payload. */
-    private static final String HTTP_USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
-            + "AppleWebKit/537.36 (KHTML, like Gecko) Chrome/60.0.3112.113 Safari/537.36";
-
     private TextView mTextViewBuild;
     private Spinner mSpinnerConfigs;
     private TextView mTextViewConfigsDirHint;
@@ -73,18 +62,9 @@ public class MainActivity extends Activity {
     private Button mButtonSwitchSlot;
 
     private List<UpdateConfig> mConfigs;
-    private AtomicInteger mUpdateEngineStatus =
-            new AtomicInteger(UpdateEngine.UpdateStatusConstants.IDLE);
-    private PayloadSpec mLastPayloadSpec;
-    private AtomicBoolean mManualSwitchSlotRequired = new AtomicBoolean(true);
-    private final PayloadSpecs mPayloadSpecs = new PayloadSpecs();
-
-    /**
-     * Listen to {@code update_engine} events.
-     */
-    private UpdateEngineCallbackImpl mUpdateEngineCallback = new UpdateEngineCallbackImpl();
 
-    private final UpdateEngine mUpdateEngine = new UpdateEngine();
+    private final UpdateManager mUpdateManager =
+            new UpdateManager(new UpdateEngine(), new PayloadSpecs());
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
@@ -109,15 +89,31 @@ public class MainActivity extends Activity {
         uiReset();
         loadUpdateConfigs();
 
-        this.mUpdateEngine.bind(mUpdateEngineCallback);
+        this.mUpdateManager.setOnEngineStatusUpdateCallback(this::onStatusUpdate);
+        this.mUpdateManager.setOnProgressUpdateCallback(this::onProgressUpdate);
+        this.mUpdateManager.setOnEngineCompleteCallback(this::onPayloadApplicationComplete);
     }
 
     @Override
     protected void onDestroy() {
-        this.mUpdateEngine.unbind();
+        this.mUpdateManager.setOnEngineStatusUpdateCallback(null);
+        this.mUpdateManager.setOnProgressUpdateCallback(null);
+        this.mUpdateManager.setOnEngineCompleteCallback(null);
         super.onDestroy();
     }
 
+    @Override
+    protected void onResume() {
+        super.onResume();
+        this.mUpdateManager.bind();
+    }
+
+    @Override
+    protected void onPause() {
+        this.mUpdateManager.unbind();
+        super.onPause();
+    }
+
     /**
      * reload button is clicked
      */
@@ -147,7 +143,7 @@ public class MainActivity extends Activity {
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
                     uiSetUpdating();
-                    applyUpdate(getSelectedConfig());
+                    mUpdateManager.applyUpdate(this, getSelectedConfig());
                 })
                 .setNegativeButton(android.R.string.cancel, null)
                 .show();
@@ -162,7 +158,7 @@ public class MainActivity extends Activity {
                 .setMessage("Do you really want to cancel running update?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    stopRunningUpdate();
+                    mUpdateManager.cancelRunningUpdate();
                 })
                 .setNegativeButton(android.R.string.cancel, null).show();
     }
@@ -177,7 +173,7 @@ public class MainActivity extends Activity {
                         + " and restore old version?")
                 .setIcon(android.R.drawable.ic_dialog_alert)
                 .setPositiveButton(android.R.string.ok, (dialog, whichButton) -> {
-                    resetUpdate();
+                    mUpdateManager.resetUpdate();
                 })
                 .setNegativeButton(android.R.string.cancel, null).show();
     }
@@ -186,7 +182,7 @@ public class MainActivity extends Activity {
      * switch slot button clicked
      */
     public void onSwitchSlotClick(View view) {
-        setSwitchSlotOnReboot();
+        mUpdateManager.setSwitchSlotOnReboot();
     }
 
     /**
@@ -194,26 +190,26 @@ public class MainActivity extends Activity {
      * be one of the values from {@link UpdateEngine.UpdateStatusConstants},
      * and {@code percent} will be from {@code 0.0} to {@code 1.0}.
      */
-    private void onStatusUpdate(int status, float percent) {
-        mProgressBar.setProgress((int) (100 * percent));
-        if (mUpdateEngineStatus.get() != status) {
-            mUpdateEngineStatus.set(status);
-            runOnUiThread(() -> {
-                Log.e("UpdateEngine", "StatusUpdate - status="
-                        + UpdateEngineStatuses.getStatusText(status)
-                        + "/" + status);
-                Toast.makeText(this, "Update Status changed", Toast.LENGTH_LONG)
-                        .show();
-                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();
-                }
-                setUiStatus(status);
-            });
-        }
+    private void onStatusUpdate(int status) {
+        runOnUiThread(() -> {
+            Log.e("UpdateEngine", "StatusUpdate - status="
+                    + UpdateEngineStatuses.getStatusText(status)
+                    + "/" + status);
+            Toast.makeText(this, "Update Status changed", Toast.LENGTH_LONG)
+                    .show();
+            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();
+            }
+            setUiStatus(status);
+        });
+    }
+
+    private void onProgressUpdate(double progress) {
+        mProgressBar.setProgress((int) (100 * progress));
     }
 
     /**
@@ -234,7 +230,7 @@ public class MainActivity extends Activity {
             setUiCompletion(errorCode);
             if (errorCode == UpdateEngineErrorCodes.UPDATED_BUT_NOT_ACTIVE) {
                 // if update was successfully applied.
-                if (mManualSwitchSlotRequired.get()) {
+                if (mUpdateManager.manualSwitchSlotRequired()) {
                     // Show "Switch Slot" button.
                     uiShowSwitchSlotInfo();
                 }
@@ -321,143 +317,4 @@ public class MainActivity extends Activity {
         return mConfigs.get(mSpinnerConfigs.getSelectedItemPosition());
     }
 
-    /**
-     * Applies the given update
-     */
-    private void applyUpdate(final UpdateConfig config) {
-        List<String> extraProperties = new ArrayList<>();
-
-        if (!config.getAbConfig().getForceSwitchSlot()) {
-            // Disable switch slot on reboot, which is enabled by default.
-            // User will enable it manually by clicking "Switch Slot" button on the screen.
-            extraProperties.add(UpdateEngineProperties.PROPERTY_DISABLE_SWITCH_SLOT_ON_REBOOT);
-            mManualSwitchSlotRequired.set(true);
-        } else {
-            mManualSwitchSlotRequired.set(false);
-        }
-
-        if (config.getInstallType() == UpdateConfig.AB_INSTALL_TYPE_NON_STREAMING) {
-            PayloadSpec payload;
-            try {
-                payload = mPayloadSpecs.forNonStreaming(config.getUpdatePackageFile());
-            } catch (IOException e) {
-                Log.e(TAG, "Error creating payload spec", e);
-                Toast.makeText(this, "Error creating payload spec", Toast.LENGTH_LONG)
-                        .show();
-                return;
-            }
-            updateEngineApplyPayload(payload, extraProperties);
-        } else {
-            Log.d(TAG, "Starting PrepareStreamingService");
-            PrepareStreamingService.startService(this, config, (code, payloadSpec) -> {
-                if (code == PrepareStreamingService.RESULT_CODE_SUCCESS) {
-                    extraProperties.add("USER_AGENT=" + HTTP_USER_AGENT);
-                    config.getStreamingMetadata()
-                            .getAuthorization()
-                            .ifPresent(s -> extraProperties.add("AUTHORIZATION=" + s));
-                    updateEngineApplyPayload(payloadSpec, extraProperties);
-                } else {
-                    Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
-                    Toast.makeText(
-                            MainActivity.this,
-                            "PrepareStreamingService failed, result code is " + code,
-                            Toast.LENGTH_LONG).show();
-                }
-            });
-        }
-    }
-
-    /**
-     * Applies given payload.
-     *
-     * UpdateEngine works asynchronously. This method doesn't wait until
-     * end of the update.
-     *
-     * @param payloadSpec contains url, offset and size to {@code PAYLOAD_BINARY_FILE_NAME}
-     * @param extraProperties additional properties to pass to {@link UpdateEngine#applyPayload}
-     */
-    private void updateEngineApplyPayload(PayloadSpec payloadSpec, List<String> extraProperties) {
-        mLastPayloadSpec = payloadSpec;
-
-        ArrayList<String> properties = new ArrayList<>(payloadSpec.getProperties());
-        if (extraProperties != null) {
-            properties.addAll(extraProperties);
-        }
-        try {
-            mUpdateEngine.applyPayload(
-                    payloadSpec.getUrl(),
-                    payloadSpec.getOffset(),
-                    payloadSpec.getSize(),
-                    properties.toArray(new String[0]));
-        } catch (Exception e) {
-            Log.e(TAG, "UpdateEngine failed to apply the update", e);
-            Toast.makeText(
-                    this,
-                    "UpdateEngine failed to apply the update",
-                    Toast.LENGTH_LONG).show();
-        }
-    }
-
-    /**
-     * Sets the new slot that has the updated partitions as the active slot,
-     * which device will boot into next time.
-     * This method is only supposed to be called after the payload is applied.
-     *
-     * Invoking {@link UpdateEngine#applyPayload} with the same payload url, offset, size
-     * and payload metadata headers doesn't trigger new update. It can be used to just switch
-     * active A/B slot.
-     *
-     * {@link UpdateEngine#applyPayload} might take several seconds to finish, and it will
-     * invoke callbacks {@link this#onStatusUpdate} and {@link this#onPayloadApplicationComplete)}.
-     */
-    private void setSwitchSlotOnReboot() {
-        Log.d(TAG, "setSwitchSlotOnReboot invoked");
-        List<String> extraProperties = new ArrayList<>();
-        // PROPERTY_SKIP_POST_INSTALL should be passed on to skip post-installation hooks.
-        extraProperties.add(UpdateEngineProperties.PROPERTY_SKIP_POST_INSTALL);
-        // It sets property SWITCH_SLOT_ON_REBOOT=1 by default.
-        // HTTP headers are not required, UpdateEngine is not expected to stream payload.
-        updateEngineApplyPayload(mLastPayloadSpec, extraProperties);
-        uiHideSwitchSlotInfo();
-    }
-
-    /**
-     * Requests update engine to stop any ongoing update. If an update has been applied,
-     * leave it as is.
-     */
-    private void stopRunningUpdate() {
-        try {
-            mUpdateEngine.cancel();
-        } catch (Exception e) {
-            Log.w(TAG, "UpdateEngine failed to stop the ongoing update", e);
-        }
-    }
-
-    /**
-     * Resets update engine to IDLE state. Requests to cancel any onging update, or to revert if an
-     * update has been applied.
-     */
-    private void resetUpdate() {
-        try {
-            mUpdateEngine.resetStatus();
-        } catch (Exception e) {
-            Log.w(TAG, "UpdateEngine failed to reset the update", e);
-        }
-    }
-
-    /**
-     * Helper class to delegate {@code update_engine} callbacks to MainActivity
-     */
-    class UpdateEngineCallbackImpl extends UpdateEngineCallback {
-        @Override
-        public void onStatusUpdate(int status, float percent) {
-            MainActivity.this.onStatusUpdate(status, percent);
-        }
-
-        @Override
-        public void onPayloadApplicationComplete(int errorCode) {
-            MainActivity.this.onPayloadApplicationComplete(errorCode);
-        }
-    }
-
 }
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/PayloadSpecs.java b/updater_sample/src/com/example/android/systemupdatersample/util/PayloadSpecs.java
index b98b97c37..f06231726 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/util/PayloadSpecs.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/util/PayloadSpecs.java
@@ -32,7 +32,9 @@ import java.util.zip.ZipEntry;
 import java.util.zip.ZipFile;
 
 /** The helper class that creates {@link PayloadSpec}. */
-public final class PayloadSpecs {
+public class PayloadSpecs {
+
+    public PayloadSpecs() {}
 
     /**
      * The payload PAYLOAD_ENTRY is stored in the zip package to comply with the Android OTA package
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineErrorCodes.java b/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineErrorCodes.java
index 6d319c5af..f06ddf7fc 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineErrorCodes.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/util/UpdateEngineErrorCodes.java
@@ -34,6 +34,7 @@ public final class UpdateEngineErrorCodes {
     * Error code from the update engine. Values must agree with the ones in
     * system/update_engine/common/error_code.h.
     */
+    public static final int UNKNOWN = -1;
     public static final int UPDATED_BUT_NOT_ACTIVE = 52;
 
     private static final SparseArray<String> CODE_TO_NAME_MAP = new SparseArray<>();
diff --git a/updater_sample/tests/Android.mk b/updater_sample/tests/Android.mk
index a1a4664dc..9aec372e3 100644
--- a/updater_sample/tests/Android.mk
+++ b/updater_sample/tests/Android.mk
@@ -23,9 +23,9 @@ LOCAL_MODULE_TAGS := tests
 LOCAL_JAVA_LIBRARIES := \
     android.test.base.stubs \
     android.test.runner.stubs \
-    guava \
+    guava
+LOCAL_STATIC_JAVA_LIBRARIES := android-support-test \
     mockito-target-minus-junit4
-LOCAL_STATIC_JAVA_LIBRARIES := android-support-test
 LOCAL_INSTRUMENTATION_FOR := SystemUpdaterSample
 LOCAL_PROGUARD_ENABLED := disabled
 
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
new file mode 100644
index 000000000..0657a5eb6
--- /dev/null
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
@@ -0,0 +1,92 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.example.android.systemupdatersample;
+
+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.os.UpdateEngine;
+import android.os.UpdateEngineCallback;
+import android.support.test.filters.SmallTest;
+import android.support.test.runner.AndroidJUnit4;
+
+import com.example.android.systemupdatersample.util.PayloadSpecs;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.junit.MockitoJUnit;
+import org.mockito.junit.MockitoRule;
+
+import java.util.function.IntConsumer;
+
+/**
+ * Tests for {@link UpdateManager}
+ */
+@RunWith(AndroidJUnit4.class)
+@SmallTest
+public class UpdateManagerTest {
+
+    @Rule
+    public MockitoRule mockito = MockitoJUnit.rule();
+
+    @Mock
+    private UpdateEngine mUpdateEngine;
+    @Mock
+    private PayloadSpecs mPayloadSpecs;
+    private UpdateManager mUpdateManager;
+
+    @Before
+    public void setUp() {
+        mUpdateManager = new UpdateManager(mUpdateEngine, mPayloadSpecs);
+    }
+
+    @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.
+        when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
+            UpdateEngineCallback callback = answer.getArgument(0);
+            callback.onStatusUpdate(/*engineStatus*/ 4, /*engineProgress*/ 0.2f);
+            return null;
+        });
+
+        mUpdateManager.setOnEngineStatusUpdateCallback(statusUpdateCallback);
+
+        // Making sure that manager.getProgress() returns correct progress
+        // in "onEngineStatusUpdate" callback.
+        doAnswer(answer -> {
+            assertEquals(0.2f, mUpdateManager.getProgress(), 1E-5);
+            return null;
+        }).when(statusUpdateCallback).accept(anyInt());
+
+        mUpdateManager.bind();
+
+        verify(statusUpdateCallback, times(1)).accept(4);
+    }
+
+}
diff --git a/updater_sample/tests/src/com/example/android/systemupdatersample/ui/MainActivityTest.java b/updater_sample/tests/src/com/example/android/systemupdatersample/ui/MainActivityTest.java
deleted file mode 100644
index 01014168a..000000000
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/ui/MainActivityTest.java
+++ /dev/null
@@ -1,48 +0,0 @@
-/*
- * Copyright (C) 2018 The Android Open Source Project
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *      http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package com.example.android.systemupdatersample.ui;
-
-import static org.junit.Assert.assertNotNull;
-
-import android.support.test.filters.MediumTest;
-import android.support.test.rule.ActivityTestRule;
-import android.support.test.runner.AndroidJUnit4;
-
-import org.junit.Rule;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-
-/**
- * Make sure that the main launcher activity opens up properly, which will be
- * verified by {@link #activityLaunches}.
- */
-@RunWith(AndroidJUnit4.class)
-@MediumTest
-public class MainActivityTest {
-
-    @Rule
-    public final ActivityTestRule<MainActivity> mActivityRule =
-            new ActivityTestRule<>(MainActivity.class);
-
-    /**
-     * Verifies that the activity under test can be launched.
-     */
-    @Test
-    public void activityLaunches() {
-        assertNotNull(mActivityRule.getActivity());
-    }
-}