Test for userdata restore in staged rollback test.

Bug: 124044231
Test: frameworks/base/tests/RollbackTest$ atest
Change-Id: Ic323b2e45c05ac4706d96fe15b109fdceaf53da4

3 files changed, 57 insertions, 50 deletions

diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
index 7be83edd91b9..a6054e8f41d0 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTest.java
@@ -19,19 +19,17 @@ package com.android.tests.rollback;
 import static com.android.tests.rollback.RollbackTestUtils.assertPackageRollbackInfoEquals;
 import static com.android.tests.rollback.RollbackTestUtils.assertRollbackInfoEquals;
 import static com.android.tests.rollback.RollbackTestUtils.getUniqueRollbackInfoForPackage;
+import static com.android.tests.rollback.RollbackTestUtils.processUserData;
 
 import android.Manifest;
 import android.app.ActivityManager;
 import android.content.BroadcastReceiver;
-import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
 import android.content.pm.VersionedPackage;
 import android.content.rollback.RollbackInfo;
 import android.content.rollback.RollbackManager;
-import android.os.Handler;
-import android.os.HandlerThread;
 import android.provider.DeviceConfig;
 import android.support.test.InstrumentationRegistry;
 import android.util.Log;
@@ -47,7 +45,6 @@ import org.junit.runners.JUnit4;
 
 import java.util.Collections;
 import java.util.concurrent.BlockingQueue;
-import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.SynchronousQueue;
 import java.util.concurrent.TimeUnit;
 
@@ -427,52 +424,6 @@ public class RollbackTest {
         }
     }
 
-    private static final String NO_RESPONSE = "NO RESPONSE";
-
-    // Calls into the test app to process user data.
-    // Asserts if the user data could not be processed or was version
-    // incompatible with the previously processed user data.
-    private void processUserData(String packageName) throws Exception {
-        Intent intent = new Intent();
-        intent.setComponent(new ComponentName(packageName,
-                    "com.android.tests.rollback.testapp.ProcessUserData"));
-        Context context = InstrumentationRegistry.getContext();
-
-        HandlerThread handlerThread = new HandlerThread("RollbackTestHandlerThread");
-        handlerThread.start();
-
-        // It can sometimes take a while after rollback before the app will
-        // receive this broadcast, so try a few times in a loop.
-        String result = NO_RESPONSE;
-        for (int i = 0; result.equals(NO_RESPONSE) && i < 5; ++i) {
-            BlockingQueue<String> resultQueue = new LinkedBlockingQueue<>();
-            context.sendOrderedBroadcast(intent, null, new BroadcastReceiver() {
-                @Override
-                public void onReceive(Context context, Intent intent) {
-                    if (getResultCode() == 1) {
-                        resultQueue.add("OK");
-                    } else {
-                        // If the test app doesn't receive the broadcast or
-                        // fails to set the result data, then getResultData
-                        // here returns the initial NO_RESPONSE data passed to
-                        // the sendOrderedBroadcast call.
-                        resultQueue.add(getResultData());
-                    }
-                }
-            }, new Handler(handlerThread.getLooper()), 0, NO_RESPONSE, null);
-
-            result = resultQueue.poll(10, TimeUnit.SECONDS);
-            if (result == null) {
-                result = "ProcessUserData broadcast timed out";
-            }
-        }
-
-        handlerThread.quit();
-        if (!"OK".equals(result)) {
-            fail(result);
-        }
-    }
-
     /**
      * Test that app user data is rolled back.
      */
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTestUtils.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTestUtils.java
index f28714c87125..2f989f316ea6 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTestUtils.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/RollbackTestUtils.java
@@ -17,6 +17,7 @@
 package com.android.tests.rollback;
 
 import android.content.BroadcastReceiver;
+import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
 import android.content.IntentFilter;
@@ -27,12 +28,15 @@ import android.content.pm.VersionedPackage;
 import android.content.rollback.PackageRollbackInfo;
 import android.content.rollback.RollbackInfo;
 import android.content.rollback.RollbackManager;
+import android.os.Handler;
+import android.os.HandlerThread;
 import android.support.test.InstrumentationRegistry;
 import android.util.Log;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -393,4 +397,53 @@ class RollbackTestUtils {
             throw new AssertionError(e);
         }
     }
+
+    private static final String NO_RESPONSE = "NO RESPONSE";
+
+    /**
+     * Calls into the test app to process user data.
+     * Asserts if the user data could not be processed or was version
+     * incompatible with the previously processed user data.
+     */
+    static void processUserData(String packageName) {
+        Intent intent = new Intent();
+        intent.setComponent(new ComponentName(packageName,
+                    "com.android.tests.rollback.testapp.ProcessUserData"));
+        Context context = InstrumentationRegistry.getContext();
+
+        HandlerThread handlerThread = new HandlerThread("RollbackTestHandlerThread");
+        handlerThread.start();
+
+        // It can sometimes take a while after rollback before the app will
+        // receive this broadcast, so try a few times in a loop.
+        String result = NO_RESPONSE;
+        for (int i = 0; result.equals(NO_RESPONSE) && i < 5; ++i) {
+            BlockingQueue<String> resultQueue = new LinkedBlockingQueue<>();
+            context.sendOrderedBroadcast(intent, null, new BroadcastReceiver() {
+                @Override
+                public void onReceive(Context context, Intent intent) {
+                    if (getResultCode() == 1) {
+                        resultQueue.add("OK");
+                    } else {
+                        // If the test app doesn't receive the broadcast or
+                        // fails to set the result data, then getResultData
+                        // here returns the initial NO_RESPONSE data passed to
+                        // the sendOrderedBroadcast call.
+                        resultQueue.add(getResultData());
+                    }
+                }
+            }, new Handler(handlerThread.getLooper()), 0, NO_RESPONSE, null);
+
+            try {
+                result = resultQueue.take();
+            } catch (InterruptedException e) {
+                throw new AssertionError(e);
+            }
+        }
+
+        handlerThread.quit();
+        if (!"OK".equals(result)) {
+            fail(result);
+        }
+    }
 }
diff --git a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
index b65917b2ae5b..59ae8d9deaeb 100644
--- a/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
+++ b/tests/RollbackTest/RollbackTest/src/com/android/tests/rollback/StagedRollbackTest.java
@@ -84,6 +84,7 @@ public class StagedRollbackTest {
 
         RollbackTestUtils.install("RollbackTestAppAv1.apk", false);
         assertEquals(1, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
+        RollbackTestUtils.processUserData(TEST_APP_A);
 
         RollbackTestUtils.installStaged(true, "RollbackTestAppAv2.apk");
 
@@ -98,6 +99,7 @@ public class StagedRollbackTest {
     @Test
     public void testApkOnlyCommitRollback() throws Exception {
         assertEquals(2, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
+        RollbackTestUtils.processUserData(TEST_APP_A);
 
         RollbackManager rm = RollbackTestUtils.getRollbackManager();
         RollbackInfo rollback = getUniqueRollbackInfoForPackage(
@@ -129,6 +131,7 @@ public class StagedRollbackTest {
     @Test
     public void testApkOnlyConfirmRollback() throws Exception {
         assertEquals(1, RollbackTestUtils.getInstalledVersion(TEST_APP_A));
+        RollbackTestUtils.processUserData(TEST_APP_A);
 
         RollbackManager rm = RollbackTestUtils.getRollbackManager();
         RollbackInfo rollback = getUniqueRollbackInfoForPackage(