8 files changed, 111 insertions, 113 deletions

diff --git a/updater_sample/AndroidManifest.xml b/updater_sample/AndroidManifest.xml
index 18d8425e1..0a2511617 100644
--- a/updater_sample/AndroidManifest.xml
+++ b/updater_sample/AndroidManifest.xml
@@ -33,7 +33,7 @@
                 <category android:name="android.intent.category.LAUNCHER" />
             </intent-filter>
         </activity>
-        <service android:name=".services.PrepareStreamingService"/>
+        <service android:name=".services.PrepareUpdateService"/>
     </application>
 
 </manifest>
diff --git a/updater_sample/README.md b/updater_sample/README.md
index f9c3fb8ec..5894cf8cd 100644
--- a/updater_sample/README.md
+++ b/updater_sample/README.md
@@ -220,7 +220,7 @@ privileged system app, so it's granted the required permissions to access
 - [x] Add Sample app update state (separate from update_engine status)
 - [x] Add smart update completion detection using onStatusUpdate
 - [x] Add pause/resume demo
-- [x] Verify system partition checksum for package
+- [-] Verify system partition checksum for package
 
 
 ## Running tests
@@ -235,8 +235,8 @@ privileged system app, so it's granted the required permissions to access
 5. Run a test file
    ```
    adb shell am instrument \
-     -w com.example.android.systemupdatersample.tests/android.support.test.runner.AndroidJUnitRunner \
-     -c com.example.android.systemupdatersample.util.PayloadSpecsTest
+     -w -e class com.example.android.systemupdatersample.UpdateManagerTest#applyUpdate_appliesPayloadToUpdateEngine \
+     com.example.android.systemupdatersample.tests/android.support.test.runner.AndroidJUnitRunner
    ```
 
 
diff --git a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
index 12a8f3f5f..c02e60846 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/UpdateManager.java
@@ -17,19 +17,18 @@
 package com.example.android.systemupdatersample;
 
 import android.content.Context;
+import android.os.Handler;
 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.services.PrepareUpdateService;
 import com.example.android.systemupdatersample.util.UpdateEngineErrorCodes;
 import com.example.android.systemupdatersample.util.UpdateEngineProperties;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.AtomicDouble;
 
-import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -50,11 +49,10 @@ 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) "
+    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);
@@ -84,9 +82,15 @@ public class UpdateManager {
     private final UpdateManager.UpdateEngineCallbackImpl
             mUpdateEngineCallback = new UpdateManager.UpdateEngineCallbackImpl();
 
-    public UpdateManager(UpdateEngine updateEngine, PayloadSpecs payloadSpecs) {
+    private final Handler mHandler;
+
+    /**
+     * @param updateEngine UpdateEngine instance.
+     * @param handler      Handler for {@link PrepareUpdateService} intent service.
+     */
+    public UpdateManager(UpdateEngine updateEngine, Handler handler) {
         this.mUpdateEngine = updateEngine;
-        this.mPayloadSpecs = payloadSpecs;
+        this.mHandler = handler;
     }
 
     /**
@@ -293,45 +297,17 @@ public class UpdateManager {
             mManualSwitchSlotRequired.set(false);
         }
 
-        if (config.getInstallType() == UpdateConfig.AB_INSTALL_TYPE_NON_STREAMING) {
-            applyAbNonStreamingUpdate(config);
-        } else {
-            applyAbStreamingUpdate(context, config);
-        }
-    }
-
-    private void applyAbNonStreamingUpdate(UpdateConfig config)
-            throws UpdaterState.InvalidTransitionException {
-        UpdateData.Builder builder = UpdateData.builder()
-                .setExtraProperties(prepareExtraProperties(config));
-
-        try {
-            builder.setPayload(mPayloadSpecs.forNonStreaming(config.getUpdatePackageFile()));
-        } catch (IOException e) {
-            Log.e(TAG, "Error creating payload spec", e);
-            setUpdaterState(UpdaterState.ERROR);
-            return;
-        }
-        updateEngineApplyPayload(builder.build());
-    }
-
-    private void applyAbStreamingUpdate(Context context, UpdateConfig config) {
-        UpdateData.Builder builder = UpdateData.builder()
-                .setExtraProperties(prepareExtraProperties(config));
-
-        Log.d(TAG, "Starting PrepareStreamingService");
-        PrepareStreamingService.startService(context, config, (code, payloadSpec) -> {
-            if (code == PrepareStreamingService.RESULT_CODE_SUCCESS) {
-                builder.setPayload(payloadSpec);
-                builder.addExtraProperty("USER_AGENT=" + HTTP_USER_AGENT);
-                config.getAbConfig()
-                        .getAuthorization()
-                        .ifPresent(s -> builder.addExtraProperty("AUTHORIZATION=" + s));
-                updateEngineApplyPayload(builder.build());
-            } else {
-                Log.e(TAG, "PrepareStreamingService failed, result code is " + code);
+        Log.d(TAG, "Starting PrepareUpdateService");
+        PrepareUpdateService.startService(context, config, mHandler, (code, payloadSpec) -> {
+            if (code != PrepareUpdateService.RESULT_CODE_SUCCESS) {
+                Log.e(TAG, "PrepareUpdateService failed, result code is " + code);
                 setUpdaterStateSilent(UpdaterState.ERROR);
+                return;
             }
+            updateEngineApplyPayload(UpdateData.builder()
+                    .setExtraProperties(prepareExtraProperties(config))
+                    .setPayload(payloadSpec)
+                    .build());
         });
     }
 
@@ -343,6 +319,12 @@ public class UpdateManager {
             // User will enable it manually by clicking "Switch Slot" button on the screen.
             extraProperties.add(UpdateEngineProperties.PROPERTY_DISABLE_SWITCH_SLOT_ON_REBOOT);
         }
+        if (config.getInstallType() == UpdateConfig.AB_INSTALL_TYPE_STREAMING) {
+            extraProperties.add("USER_AGENT=" + HTTP_USER_AGENT);
+            config.getAbConfig()
+                    .getAuthorization()
+                    .ifPresent(s -> extraProperties.add("AUTHORIZATION=" + s));
+        }
         return extraProperties;
     }
 
@@ -497,14 +479,14 @@ public class UpdateManager {
      * system/update_engine/binder_service_android.cc in
      * function BinderUpdateEngineAndroidService::bind).
      *
-     * @param status one of {@link UpdateEngine.UpdateStatusConstants}.
+     * @param status   one of {@link UpdateEngine.UpdateStatusConstants}.
      * @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));
+                "onStatusUpdate invoked, status=%s, progress=%.2f",
+                status,
+                progress));
 
         int previousStatus = mUpdateEngineStatus.get();
         mUpdateEngineStatus.set(status);
@@ -555,7 +537,6 @@ public class UpdateManager {
     }
 
     /**
-     *
      * Contains update data - PayloadSpec and extra properties list.
      *
      * <p>{@code mPayload} contains url, offset and size to {@code PAYLOAD_BINARY_FILE_NAME}.
diff --git a/updater_sample/src/com/example/android/systemupdatersample/services/PrepareStreamingService.java b/updater_sample/src/com/example/android/systemupdatersample/services/PrepareUpdateService.java
index 931404857..06581bee3 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/services/PrepareStreamingService.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/services/PrepareUpdateService.java
@@ -28,6 +28,7 @@ import android.os.Bundle;
 import android.os.Handler;
 import android.os.RecoverySystem;
 import android.os.ResultReceiver;
+import android.os.UpdateEngine;
 import android.util.Log;
 
 import com.example.android.systemupdatersample.PayloadSpec;
@@ -49,10 +50,10 @@ import java.util.Optional;
  * without downloading the whole package. And it constructs {@link PayloadSpec}.
  * All this work required to install streaming A/B updates.
  *
- * PrepareStreamingService runs on it's own thread. It will notify activity
+ * PrepareUpdateService runs on it's own thread. It will notify activity
  * using interface {@link UpdateResultCallback} when update is ready to install.
  */
-public class PrepareStreamingService extends IntentService {
+public class PrepareUpdateService extends IntentService {
 
     /**
      * UpdateResultCallback result codes.
@@ -61,62 +62,63 @@ public class PrepareStreamingService extends IntentService {
     public static final int RESULT_CODE_ERROR = 1;
 
     /**
-     * This interface is used to send results from {@link PrepareStreamingService} to
+     * Extra params that will be sent to IntentService.
+     */
+    public static final String EXTRA_PARAM_CONFIG = "config";
+    public static final String EXTRA_PARAM_RESULT_RECEIVER = "result-receiver";
+
+    /**
+     * This interface is used to send results from {@link PrepareUpdateService} to
      * {@code MainActivity}.
      */
     public interface UpdateResultCallback {
-
         /**
          * Invoked when files are downloaded and payload spec is constructed.
          *
-         * @param resultCode result code, values are defined in {@link PrepareStreamingService}
+         * @param resultCode  result code, values are defined in {@link PrepareUpdateService}
          * @param payloadSpec prepared payload spec for streaming update
          */
         void onReceiveResult(int resultCode, PayloadSpec payloadSpec);
     }
 
     /**
-     * Starts PrepareStreamingService.
+     * Starts PrepareUpdateService.
      *
-     * @param context application context
-     * @param config update config
+     * @param context        application context
+     * @param config         update config
      * @param resultCallback callback that will be called when the update is ready to be installed
      */
     public static void startService(Context context,
             UpdateConfig config,
+            Handler handler,
             UpdateResultCallback resultCallback) {
-        Log.d(TAG, "Starting PrepareStreamingService");
-        ResultReceiver receiver = new CallbackResultReceiver(new Handler(), resultCallback);
-        Intent intent = new Intent(context, PrepareStreamingService.class);
+        Log.d(TAG, "Starting PrepareUpdateService");
+        ResultReceiver receiver = new CallbackResultReceiver(handler, resultCallback);
+        Intent intent = new Intent(context, PrepareUpdateService.class);
         intent.putExtra(EXTRA_PARAM_CONFIG, config);
         intent.putExtra(EXTRA_PARAM_RESULT_RECEIVER, receiver);
         context.startService(intent);
     }
 
-    public PrepareStreamingService() {
+    public PrepareUpdateService() {
         super(TAG);
     }
 
-    private static final String TAG = "PrepareStreamingService";
-
-    /**
-     * Extra params that will be sent from Activity to IntentService.
-     */
-    private static final String EXTRA_PARAM_CONFIG = "config";
-    private static final String EXTRA_PARAM_RESULT_RECEIVER = "result-receiver";
+    private static final String TAG = "PrepareUpdateService";
 
     /**
      * The files that should be downloaded before streaming.
      */
     private static final ImmutableSet<String> PRE_STREAMING_FILES_SET =
             ImmutableSet.of(
-                PackageFiles.CARE_MAP_FILE_NAME,
-                PackageFiles.COMPATIBILITY_ZIP_FILE_NAME,
-                PackageFiles.METADATA_FILE_NAME,
-                PackageFiles.PAYLOAD_PROPERTIES_FILE_NAME
+                    PackageFiles.CARE_MAP_FILE_NAME,
+                    PackageFiles.COMPATIBILITY_ZIP_FILE_NAME,
+                    PackageFiles.METADATA_FILE_NAME,
+                    PackageFiles.PAYLOAD_PROPERTIES_FILE_NAME
             );
 
     private final PayloadSpecs mPayloadSpecs = new PayloadSpecs();
+    private final UpdateEngine mUpdateEngine = new UpdateEngine();
 
     @Override
     protected void onHandleIntent(Intent intent) {
@@ -142,6 +144,10 @@ public class PrepareStreamingService extends IntentService {
     private PayloadSpec execute(UpdateConfig config)
             throws IOException, PreparationFailedException {
 
+        if (config.getInstallType() == UpdateConfig.AB_INSTALL_TYPE_NON_STREAMING) {
+            return mPayloadSpecs.forNonStreaming(config.getUpdatePackageFile());
+        }
+
         downloadPreStreamingFiles(config, OTA_PACKAGE_DIR);
 
         Optional<UpdateConfig.PackageFile> payloadBinary =
@@ -176,6 +182,7 @@ public class PrepareStreamingService extends IntentService {
      * Downloads files defined in {@link UpdateConfig#getAbConfig()}
      * and exists in {@code PRE_STREAMING_FILES_SET}, and put them
      * in directory {@code dir}.
+     *
      * @throws IOException when can't download a file
      */
     private void downloadPreStreamingFiles(UpdateConfig config, String dir)
@@ -212,7 +219,7 @@ public class PrepareStreamingService extends IntentService {
     }
 
     /**
-     * Used by {@link PrepareStreamingService} to pass {@link PayloadSpec}
+     * Used by {@link PrepareUpdateService} to pass {@link PayloadSpec}
      * to {@link UpdateResultCallback#onReceiveResult}.
      */
     private static class CallbackResultReceiver extends ResultReceiver {
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 fc9fddd70..6d1e4c35a 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/ui/MainActivity.java
@@ -21,6 +21,7 @@ import android.app.AlertDialog;
 import android.graphics.Color;
 import android.os.Build;
 import android.os.Bundle;
+import android.os.Handler;
 import android.os.UpdateEngine;
 import android.util.Log;
 import android.view.View;
@@ -34,7 +35,6 @@ import com.example.android.systemupdatersample.R;
 import com.example.android.systemupdatersample.UpdateConfig;
 import com.example.android.systemupdatersample.UpdateManager;
 import com.example.android.systemupdatersample.UpdaterState;
-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.UpdateEngineStatuses;
@@ -67,7 +67,7 @@ public class MainActivity extends Activity {
     private List<UpdateConfig> mConfigs;
 
     private final UpdateManager mUpdateManager =
-            new UpdateManager(new UpdateEngine(), new PayloadSpecs());
+            new UpdateManager(new UpdateEngine(), new Handler());
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
diff --git a/updater_sample/src/com/example/android/systemupdatersample/util/FileDownloader.java b/updater_sample/src/com/example/android/systemupdatersample/util/FileDownloader.java
index ddd0919b8..0f9083d27 100644
--- a/updater_sample/src/com/example/android/systemupdatersample/util/FileDownloader.java
+++ b/updater_sample/src/com/example/android/systemupdatersample/util/FileDownloader.java
@@ -30,7 +30,7 @@ import java.net.URLConnection;
  * Downloads chunk of a file from given url using {@code offset} and {@code size},
  * and saves to a given location.
  *
- * In real-life application this helper class should download from HTTP Server,
+ * In a real-life application this helper class should download from HTTP Server,
  * but in this sample app it will only download from a local file.
  */
 public final class FileDownloader {
diff --git a/updater_sample/tests/Android.bp b/updater_sample/tests/Android.bp
index e43440518..806babd9e 100644
--- a/updater_sample/tests/Android.bp
+++ b/updater_sample/tests/Android.bp
@@ -24,6 +24,7 @@ android_test {
 
     static_libs: [
         "androidx.test.runner",
+        "androidx.test.rules",
         "mockito-target-minus-junit4",
         "guava",
     ],
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 6bef6de0e..af891f007 100644
--- a/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
+++ b/updater_sample/tests/src/com/example/android/systemupdatersample/UpdateManagerTest.java
@@ -18,21 +18,26 @@ package com.example.android.systemupdatersample;
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.ArgumentMatchers.any;
+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.content.Intent;
+import android.os.Bundle;
+import android.os.ResultReceiver;
 import android.os.UpdateEngine;
 import android.os.UpdateEngineCallback;
 
 import androidx.test.InstrumentationRegistry;
+import androidx.test.annotation.UiThreadTest;
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
 
+import com.example.android.systemupdatersample.services.PrepareUpdateService;
 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;
 
@@ -44,7 +49,6 @@ import org.mockito.Mock;
 import org.mockito.junit.MockitoJUnit;
 import org.mockito.junit.MockitoRule;
 
-import java.io.File;
 import java.io.IOException;
 import java.io.InputStreamReader;
 
@@ -61,49 +65,39 @@ public class UpdateManagerTest {
     @Mock
     private UpdateEngine mUpdateEngine;
     @Mock
-    private PayloadSpecs mPayloadSpecs;
+    private Context mMockContext;
     private UpdateManager mSubject;
-    private Context mContext;
-    private UpdateConfig mNonStreamingUpdate003;
+    private Context mTestContext;
+    private UpdateConfig mStreamingUpdate002;
 
     @Before
     public void setUp() throws Exception {
-        mContext = InstrumentationRegistry.getContext();
-        mSubject = new UpdateManager(mUpdateEngine, mPayloadSpecs);
-        mNonStreamingUpdate003 =
-                UpdateConfig.fromJson(readResource(R.raw.update_config_003_nonstream));
+        mTestContext = InstrumentationRegistry.getContext();
+        mSubject = new UpdateManager(mUpdateEngine, null);
+        mStreamingUpdate002 =
+                UpdateConfig.fromJson(readResource(R.raw.update_config_002_stream));
     }
 
     @Test
     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(
-                    UpdateEngine.UpdateStatusConstants.IDLE,
-                    /*engineProgress*/ 0.0f);
-            return null;
-        });
-
-        mSubject.bind();
-        mSubject.applyUpdate(null, mNonStreamingUpdate003);
+        mockContextStartServiceAnswer(buildMockPayloadSpec());
+        mSubject.applyUpdate(mMockContext, mStreamingUpdate002);
 
         verify(mUpdateEngine).applyPayload(
                 "file://blah",
                 120,
                 340,
-                new String[] {
-                        "SWITCH_SLOT_ON_REBOOT=0" // ab_config.force_switch_slot = false
+                new String[]{
+                        "SWITCH_SLOT_ON_REBOOT=0", // ab_config.force_switch_slot = false
+                        "USER_AGENT=" + UpdateManager.HTTP_USER_AGENT
                 });
     }
 
     @Test
-    public void stateIsRunningAndEngineStatusIsIdle_reApplyLastUpdate() throws Exception {
-        PayloadSpec payload = buildMockPayloadSpec();
-        when(mPayloadSpecs.forNonStreaming(any(File.class))).thenReturn(payload);
+    @UiThreadTest
+    public void stateIsRunningAndEngineStatusIsIdle_reApplyLastUpdate() throws Throwable {
+        mockContextStartServiceAnswer(buildMockPayloadSpec());
+        // UpdateEngine always returns IDLE status.
         when(mUpdateEngine.bind(any(UpdateEngineCallback.class))).thenAnswer(answer -> {
             // When UpdateManager is bound to update_engine, it passes
             // UpdateEngineCallback as a callback to update_engine.
@@ -115,21 +109,36 @@ public class UpdateManagerTest {
         });
 
         mSubject.bind();
-        mSubject.applyUpdate(null, mNonStreamingUpdate003);
+        mSubject.applyUpdate(mMockContext, mStreamingUpdate002);
         mSubject.unbind();
         mSubject.bind(); // re-bind - now it should re-apply last update
 
         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
+                new String[]{
+                        "SWITCH_SLOT_ON_REBOOT=0", // ab_config.force_switch_slot = false
+                        "USER_AGENT=" + UpdateManager.HTTP_USER_AGENT
                 });
     }
 
+    private void mockContextStartServiceAnswer(PayloadSpec payloadSpec) {
+        doAnswer(args -> {
+            Intent intent = args.getArgument(0);
+            ResultReceiver resultReceiver = intent.getParcelableExtra(
+                    PrepareUpdateService.EXTRA_PARAM_RESULT_RECEIVER);
+            Bundle b = new Bundle();
+            b.putSerializable(
+                    /* PrepareUpdateService.CallbackResultReceiver.BUNDLE_PARAM_PAYLOAD_SPEC */
+                    "payload-spec",
+                    payloadSpec);
+            resultReceiver.send(PrepareUpdateService.RESULT_CODE_SUCCESS, b);
+            return null;
+        }).when(mMockContext).startService(any(Intent.class));
+    }
+
     private PayloadSpec buildMockPayloadSpec() {
         PayloadSpec payload = mock(PayloadSpec.class);
         when(payload.getUrl()).thenReturn("file://blah");
@@ -141,7 +150,7 @@ public class UpdateManagerTest {
 
     private String readResource(int id) throws IOException {
         return CharStreams.toString(new InputStreamReader(
-                mContext.getResources().openRawResource(id)));
+                mTestContext.getResources().openRawResource(id)));
     }
 
 }