Merge "Add Notification bubbles setting to the validators to fix tests"
diff --git a/Android.bp b/Android.bp
index 121decb..ee3cb8f 100644
--- a/Android.bp
+++ b/Android.bp
@@ -234,6 +234,11 @@
             "wifi/java",
         ],
     },
+
+    required: [
+        // TODO: remove gps_debug when the build system propagates "required" properly.
+        "gps_debug.conf",
+    ],
 }
 
 // Collection of classes that are generated from non-Java files that are not listed in
@@ -307,11 +312,6 @@
 
     static_libs: ["framework-internal-utils"],
 
-    required: [
-        // TODO: remove gps_debug when the build system propagates "required" properly.
-        "gps_debug.conf",
-    ],
-
     dxflags: [
         "--core-library",
         "--multi-dex",
diff --git a/core/java/android/app/AppOpsManager.java b/core/java/android/app/AppOpsManager.java
index 1f283ea..bb87d96 100644
--- a/core/java/android/app/AppOpsManager.java
+++ b/core/java/android/app/AppOpsManager.java
@@ -5845,39 +5845,42 @@
                     ActivityThread.currentOpPackageName())) {
                 // This app is noting an app-op for itself. Deliver immediately.
                 sNotedAppOpsCollector.onSelfNoted(new SyncNotedAppOp(code));
-            } else if (binderUid != null && binderUid == uid) {
-                // We are inside of a two-way binder call. Delivered to caller via
-                // {@link #prefixParcelWithAppOpsIfNeeded}
-                long[] appOpsNotedInThisBinderTransaction;
 
-                appOpsNotedInThisBinderTransaction = sAppOpsNotedInThisBinderTransaction.get();
-                if (appOpsNotedInThisBinderTransaction == null) {
-                    appOpsNotedInThisBinderTransaction = new long[2];
-                    sAppOpsNotedInThisBinderTransaction.set(appOpsNotedInThisBinderTransaction);
-                }
+                return;
+            }
+        }
 
-                if (code < 64) {
-                    appOpsNotedInThisBinderTransaction[0] |= 1L << code;
-                } else {
-                    appOpsNotedInThisBinderTransaction[1] |= 1L << (code - 64);
-                }
+        if (binderUid != null && binderUid == uid) {
+            // If this is inside of a two-way binder call: Delivered to caller via
+            // {@link #prefixParcelWithAppOpsIfNeeded}
+            long[] appOpsNotedInThisBinderTransaction;
+
+            appOpsNotedInThisBinderTransaction = sAppOpsNotedInThisBinderTransaction.get();
+            if (appOpsNotedInThisBinderTransaction == null) {
+                appOpsNotedInThisBinderTransaction = new long[2];
+                sAppOpsNotedInThisBinderTransaction.set(appOpsNotedInThisBinderTransaction);
+            }
+
+            if (code < 64) {
+                appOpsNotedInThisBinderTransaction[0] |= 1L << code;
             } else {
-                // We cannot deliver the note synchronous. Hence send it to the system server to
-                // notify the noted process.
-                if (message == null) {
-                    // Default message is a stack trace
-                    message = getFormattedStackTrace();
-                }
+                appOpsNotedInThisBinderTransaction[1] |= 1L << (code - 64);
+            }
+        } else {
+            // Cannot deliver the note synchronous: Hence send it to the system server to
+            // notify the noted process.
+            if (message == null) {
+                // Default message is a stack trace
+                message = getFormattedStackTrace();
+            }
 
-                long token = Binder.clearCallingIdentity();
-                try {
-                    mService.noteAsyncOp(mContext.getOpPackageName(), uid, packageName, code,
-                            message);
-                } catch (RemoteException e) {
-                    e.rethrowFromSystemServer();
-                } finally {
-                    Binder.restoreCallingIdentity(token);
-                }
+            long token = Binder.clearCallingIdentity();
+            try {
+                mService.noteAsyncOp(mContext.getOpPackageName(), uid, packageName, code, message);
+            } catch (RemoteException e) {
+                e.rethrowFromSystemServer();
+            } finally {
+                Binder.restoreCallingIdentity(token);
             }
         }
     }
diff --git a/core/java/com/android/internal/os/ZygoteInit.java b/core/java/com/android/internal/os/ZygoteInit.java
index b5d4945..31995f7 100644
--- a/core/java/com/android/internal/os/ZygoteInit.java
+++ b/core/java/com/android/internal/os/ZygoteInit.java
@@ -460,7 +460,7 @@
         ZygoteHooks.gcAndFinalize();
     }
 
-    private static boolean profileSystemServer() {
+    private static boolean shouldProfileSystemServer() {
         boolean defaultValue = SystemProperties.getBoolean("dalvik.vm.profilesystemserver",
                 /*default=*/ false);
         // Can't use DeviceConfig since it's not initialized at this point.
@@ -492,7 +492,7 @@
             }
             // Capturing profiles is only supported for debug or eng builds since selinux normally
             // prevents it.
-            if (profileSystemServer() && (Build.IS_USERDEBUG || Build.IS_ENG)) {
+            if (shouldProfileSystemServer() && (Build.IS_USERDEBUG || Build.IS_ENG)) {
                 try {
                     Log.d(TAG, "Preparing system server profile");
                     prepareSystemServerProfile(systemServerClasspath);
@@ -765,8 +765,7 @@
             Zygote.applyDebuggerSystemProperty(parsedArgs);
             Zygote.applyInvokeWithSystemProperty(parsedArgs);
 
-            if (profileSystemServer()) {
-
+            if (shouldProfileSystemServer()) {
                 parsedArgs.mRuntimeFlags |= Zygote.PROFILE_SYSTEM_SERVER;
             }
 
diff --git a/core/proto/android/stats/docsui/docsui_enums.proto b/core/proto/android/stats/docsui/docsui_enums.proto
index 655b5e3..f648912 100644
--- a/core/proto/android/stats/docsui/docsui_enums.proto
+++ b/core/proto/android/stats/docsui/docsui_enums.proto
@@ -184,6 +184,8 @@
     TYPE_CHIP_DOCS = 4;
     TYPE_SEARCH_HISTORY = 5;
     TYPE_SEARCH_STRING = 6;
+    TYPE_CHIP_LARGE_FILES = 7;
+    TYPE_CHIP_FROM_THIS_WEEK = 8;
 }
 
 enum SearchMode {
diff --git a/packages/PackageInstaller/src/com/android/packageinstaller/PackageInstallerActivity.java b/packages/PackageInstaller/src/com/android/packageinstaller/PackageInstallerActivity.java
index 55ff591..a2fa461 100644
--- a/packages/PackageInstaller/src/com/android/packageinstaller/PackageInstallerActivity.java
+++ b/packages/PackageInstaller/src/com/android/packageinstaller/PackageInstallerActivity.java
@@ -135,6 +135,7 @@
 
         mEnableOk = true;
         mOk.setEnabled(true);
+        mOk.setFilterTouchesWhenObscured(true);
     }
 
     /**
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
index b964edf..720074b 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
@@ -217,8 +217,6 @@
         }
     }
 
-    private static KeyguardUpdateMonitor sInstance;
-
     private final Context mContext;
     private final boolean mIsPrimaryUser;
     HashMap<Integer, SimData> mSimDatas = new HashMap<Integer, SimData>();
@@ -1389,15 +1387,6 @@
         Trace.endSection();
     }
 
-
-    /** Provides access to the static instance. */
-    public static KeyguardUpdateMonitor getInstance(Context context) {
-        if (sInstance == null) {
-            sInstance = new KeyguardUpdateMonitor(context, Looper.getMainLooper());
-        }
-        return sInstance;
-    }
-
     protected void handleStartedGoingToSleep(int arg1) {
         checkIsHandlerThread();
         mLockIconPressed = false;
diff --git a/packages/SystemUI/src/com/android/systemui/SystemUIAppComponentFactory.java b/packages/SystemUI/src/com/android/systemui/SystemUIAppComponentFactory.java
index 0f99674..e8b5285 100644
--- a/packages/SystemUI/src/com/android/systemui/SystemUIAppComponentFactory.java
+++ b/packages/SystemUI/src/com/android/systemui/SystemUIAppComponentFactory.java
@@ -21,6 +21,8 @@
 import android.content.ContentProvider;
 import android.content.Context;
 import android.content.Intent;
+import android.os.Process;
+import android.util.Log;
 
 import androidx.annotation.NonNull;
 import androidx.core.app.AppComponentFactory;
@@ -38,6 +40,7 @@
  */
 public class SystemUIAppComponentFactory extends AppComponentFactory {
 
+    private static final String TAG = "AppComponentFactory";
     @Inject
     public ContextComponentHelper mComponentHelper;
 
@@ -57,6 +60,9 @@
                         SystemUIFactory.createFromConfig(context);
                         SystemUIFactory.getInstance().getRootComponent().inject(
                                 SystemUIAppComponentFactory.this);
+                        Log.d(TAG, "Initialized during Application creation in Process "
+                                + Process.myPid() + ", Thread " + Process.myTid());
+                        Log.d(TAG, "mComponentHelper: " + mComponentHelper);
                     }
             );
         }
@@ -77,6 +83,8 @@
                         SystemUIFactory.createFromConfig(context);
                         SystemUIFactory.getInstance().getRootComponent().inject(
                                 contentProvider);
+                        Log.d(TAG, "Initialized during ContentProvider creation in Process "
+                                + Process.myPid() + ", Thread " + Process.myTid());
                     }
             );
         }
@@ -89,6 +97,12 @@
     public Service instantiateServiceCompat(
             @NonNull ClassLoader cl, @NonNull String className, Intent intent)
             throws InstantiationException, IllegalAccessException, ClassNotFoundException {
+        if (mComponentHelper == null) {
+            // Everything is about to crash if this is true, but that is inevitable. We either crash
+            // here or crash lower in the stack. Better to crash early!
+            Log.wtf(TAG, "Uninitialized mComponentHelper in Process" + Process.myPid() + ", Thread "
+                    + Process.myTid());
+        }
         Service service = mComponentHelper.resolveService(className);
         if (service != null) {
             return service;
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
index 00bfb3f..bab64db 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeTriggers.java
@@ -417,6 +417,9 @@
         mDozeSensors.dump(pw);
     }
 
+    /**
+     * @see DozeSensors.ProxSensor
+     */
     private abstract class ProximityCheck implements SensorEventListener, Runnable {
         private static final int TIMEOUT_DELAY_MS = 500;
 
@@ -428,6 +431,7 @@
         private boolean mRegistered;
         private boolean mFinished;
         private float mMaxRange;
+        private boolean mUsingBrightnessSensor;
 
         protected abstract void onProximityResult(int result);
 
@@ -435,6 +439,7 @@
             Preconditions.checkState(!mFinished && !mRegistered);
             Sensor sensor = DozeSensors.findSensorWithType(mSensorManager,
                     mContext.getString(R.string.doze_brightness_sensor_type));
+            mUsingBrightnessSensor = sensor != null;
             if (sensor == null) {
                 sensor = mSensorManager.getDefaultSensor(Sensor.TYPE_PROXIMITY);
             }
@@ -453,6 +458,9 @@
             mRegistered = true;
         }
 
+        /**
+         * @see DozeSensors.ProxSensor#onSensorChanged(SensorEvent)
+         */
         @Override
         public void onSensorChanged(SensorEvent event) {
             if (event.values.length == 0) {
@@ -462,7 +470,14 @@
                 if (DozeMachine.DEBUG) {
                     Log.d(TAG, "ProxCheck: Event: value=" + event.values[0] + " max=" + mMaxRange);
                 }
-                final boolean isNear = event.values[0] < mMaxRange;
+                final boolean isNear;
+                if (mUsingBrightnessSensor) {
+                    // The custom brightness sensor is gated by the proximity sensor and will
+                    // return 0 whenever prox is covered.
+                    isNear = event.values[0] == 0;
+                } else {
+                    isNear = event.values[0] < mMaxRange;
+                }
                 finishWithResult(isNear ? RESULT_NEAR : RESULT_FAR);
             }
         }
diff --git a/services/core/java/com/android/server/rollback/Rollback.java b/services/core/java/com/android/server/rollback/Rollback.java
index 0d5746b..6769fe0 100644
--- a/services/core/java/com/android/server/rollback/Rollback.java
+++ b/services/core/java/com/android/server/rollback/Rollback.java
@@ -65,7 +65,7 @@
     /**
      * The directory where the rollback data is stored.
      */
-    public final File backupDir;
+    private final File mBackupDir;
 
     /**
      * The time when the upgrade occurred, for purposes of expiring
@@ -74,24 +74,24 @@
      * The timestamp is not applicable for all rollback states, but we make
      * sure to keep it non-null to avoid potential errors there.
      */
-    public @NonNull Instant timestamp;
+    private @NonNull Instant mTimestamp;
 
     /**
      * The session ID for the staged session if this rollback data represents a staged session,
      * {@code -1} otherwise.
      */
-    public final int stagedSessionId;
+    private final int mStagedSessionId;
 
     /**
      * The current state of the rollback.
      * ENABLING, AVAILABLE, or COMMITTED.
      */
-    public @RollbackState int state;
+    private @RollbackState int mState;
 
     /**
      * The id of the post-reboot apk session for a staged install, if any.
      */
-    public int apkSessionId = -1;
+    private int mApkSessionId = -1;
 
     /**
      * True if we are expecting the package manager to call restoreUserData
@@ -99,7 +99,7 @@
      * has not yet been fully applied.
      */
     // NOTE: All accesses to this field are from the RollbackManager handler thread.
-    public boolean restoreUserDataInProgress = false;
+    private boolean mRestoreUserDataInProgress = false;
 
     /**
      * Constructs a new, empty Rollback instance.
@@ -114,10 +114,10 @@
                 /* isStaged */ stagedSessionId != -1,
                 /* causePackages */ new ArrayList<>(),
                 /* committedSessionId */ -1);
-        this.backupDir = backupDir;
-        this.stagedSessionId = stagedSessionId;
-        this.state = ROLLBACK_STATE_ENABLING;
-        this.timestamp = Instant.now();
+        mBackupDir = backupDir;
+        mStagedSessionId = stagedSessionId;
+        mState = ROLLBACK_STATE_ENABLING;
+        mTimestamp = Instant.now();
     }
 
     /**
@@ -126,21 +126,115 @@
     Rollback(RollbackInfo info, File backupDir, Instant timestamp, int stagedSessionId,
             @RollbackState int state, int apkSessionId, boolean restoreUserDataInProgress) {
         this.info = info;
-        this.backupDir = backupDir;
-        this.timestamp = timestamp;
-        this.stagedSessionId = stagedSessionId;
-        this.state = state;
-        this.apkSessionId = apkSessionId;
-        this.restoreUserDataInProgress = restoreUserDataInProgress;
+        mBackupDir = backupDir;
+        mTimestamp = timestamp;
+        mStagedSessionId = stagedSessionId;
+        mState = state;
+        mApkSessionId = apkSessionId;
+        mRestoreUserDataInProgress = restoreUserDataInProgress;
     }
 
     /**
      * Whether the rollback is for rollback of a staged install.
      */
-    public boolean isStaged() {
+    boolean isStaged() {
         return info.isStaged();
     }
 
+    /**
+     * Returns the directory in which rollback data should be stored.
+     */
+    File getBackupDir() {
+        return mBackupDir;
+    }
+
+    /**
+     * Returns the time when the upgrade occurred, for purposes of expiring rollback data.
+     */
+    Instant getTimestamp() {
+        return mTimestamp;
+    }
+
+    /**
+     * Sets the time at which upgrade occurred.
+     */
+    void setTimestamp(Instant timestamp) {
+        mTimestamp = timestamp;
+    }
+
+    /**
+     * Returns the session ID for the staged session if this rollback data represents a staged
+     * session, or {@code -1} otherwise.
+     */
+    int getStagedSessionId() {
+        return mStagedSessionId;
+    }
+
+    /**
+     * Returns true if the rollback is in the ENABLING state.
+     */
+    boolean isEnabling() {
+        return mState == ROLLBACK_STATE_ENABLING;
+    }
+
+    /**
+     * Returns true if the rollback is in the AVAILABLE state.
+     */
+    boolean isAvailable() {
+        return mState == ROLLBACK_STATE_AVAILABLE;
+    }
+
+    /**
+     * Returns true if the rollback is in the COMMITTED state.
+     */
+    boolean isCommitted() {
+        return mState == ROLLBACK_STATE_COMMITTED;
+    }
+
+    /**
+     * Sets the state of the rollback to AVAILABLE.
+     */
+    void setAvailable() {
+        mState = ROLLBACK_STATE_AVAILABLE;
+    }
+
+    /**
+     * Sets the state of the rollback to COMMITTED.
+     */
+    void setCommitted() {
+        mState = ROLLBACK_STATE_COMMITTED;
+    }
+
+    /**
+     * Returns the id of the post-reboot apk session for a staged install, if any.
+     */
+    int getApkSessionId() {
+        return mApkSessionId;
+    }
+
+    /**
+     * Sets the id of the post-reboot apk session for a staged install.
+     */
+    void setApkSessionId(int apkSessionId) {
+        mApkSessionId = apkSessionId;
+    }
+
+    /**
+     * Returns true if we are expecting the package manager to call restoreUserData for this
+     * rollback because it has just been committed but the rollback has not yet been fully applied.
+     */
+    boolean isRestoreUserDataInProgress() {
+        return mRestoreUserDataInProgress;
+    }
+
+    /**
+     * Sets whether we are expecting the package manager to call restoreUserData for this
+     * rollback because it has just been committed but the rollback has not yet been fully applied.
+     */
+    void setRestoreUserDataInProgress(boolean restoreUserDataInProgress) {
+        mRestoreUserDataInProgress = restoreUserDataInProgress;
+    }
+
     static String rollbackStateToString(@RollbackState int state) {
         switch (state) {
             case Rollback.ROLLBACK_STATE_ENABLING: return "enabling";
@@ -160,7 +254,7 @@
         throw new ParseException("Invalid rollback state: " + state, 0);
     }
 
-    public String getStateAsString() {
-        return rollbackStateToString(state);
+    String getStateAsString() {
+        return rollbackStateToString(mState);
     }
 }
diff --git a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
index 3147bc6..96d284b 100644
--- a/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
+++ b/services/core/java/com/android/server/rollback/RollbackManagerServiceImpl.java
@@ -282,7 +282,7 @@
             List<RollbackInfo> rollbacks = new ArrayList<>();
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 Rollback rollback = mRollbacks.get(i);
-                if (rollback.state == Rollback.ROLLBACK_STATE_AVAILABLE) {
+                if (rollback.isAvailable()) {
                     rollbacks.add(rollback.info);
                 }
             }
@@ -298,7 +298,7 @@
             List<RollbackInfo> rollbacks = new ArrayList<>();
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 Rollback rollback = mRollbacks.get(i);
-                if (rollback.state == Rollback.ROLLBACK_STATE_COMMITTED) {
+                if (rollback.isCommitted()) {
                     rollbacks.add(rollback.info);
                 }
             }
@@ -332,7 +332,7 @@
                     Iterator<Rollback> iter = mRollbacks.iterator();
                     while (iter.hasNext()) {
                         Rollback rollback = iter.next();
-                        rollback.timestamp = rollback.timestamp.plusMillis(timeDifference);
+                        rollback.setTimestamp(rollback.getTimestamp().plusMillis(timeDifference));
                         saveRollback(rollback);
                     }
                 }
@@ -358,7 +358,7 @@
         Slog.i(TAG, "Initiating rollback");
 
         Rollback rollback = getRollbackForId(rollbackId);
-        if (rollback == null || rollback.state != Rollback.ROLLBACK_STATE_AVAILABLE) {
+        if (rollback == null || !rollback.isAvailable()) {
             sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE_ROLLBACK_UNAVAILABLE,
                     "Rollback unavailable");
             return;
@@ -454,8 +454,8 @@
                                     // TODO: Could this cause a rollback to be
                                     // resurrected if it should otherwise have
                                     // expired by now?
-                                    rollback.state = Rollback.ROLLBACK_STATE_AVAILABLE;
-                                    rollback.restoreUserDataInProgress = false;
+                                    rollback.setAvailable();
+                                    rollback.setRestoreUserDataInProgress(false);
                                 }
                                 sendFailure(statusReceiver, RollbackManager.STATUS_FAILURE_INSTALL,
                                         "Rollback downgrade install failed: "
@@ -468,7 +468,7 @@
                                 if (!rollback.isStaged()) {
                                     // All calls to restoreUserData should have
                                     // completed by now for a non-staged install.
-                                    rollback.restoreUserDataInProgress = false;
+                                    rollback.setRestoreUserDataInProgress(false);
                                 }
 
                                 rollback.info.setCommittedSessionId(parentSessionId);
@@ -490,8 +490,8 @@
             );
 
             synchronized (mLock) {
-                rollback.state = Rollback.ROLLBACK_STATE_COMMITTED;
-                rollback.restoreUserDataInProgress = true;
+                rollback.setCommitted();
+                rollback.setRestoreUserDataInProgress(true);
             }
             parentSession.commit(receiver.getIntentSender());
         } catch (IOException e) {
@@ -618,9 +618,9 @@
             synchronized (mLock) {
                 for (Rollback rollback : mRollbacks) {
                     if (rollback.isStaged()) {
-                        if (rollback.state == Rollback.ROLLBACK_STATE_ENABLING) {
+                        if (rollback.isEnabling()) {
                             enabling.add(rollback);
-                        } else if (rollback.restoreUserDataInProgress) {
+                        } else if (rollback.isRestoreUserDataInProgress()) {
                             restoreInProgress.add(rollback);
                         }
 
@@ -635,8 +635,8 @@
 
             for (Rollback rollback : enabling) {
                 PackageInstaller installer = mContext.getPackageManager().getPackageInstaller();
-                PackageInstaller.SessionInfo session = installer.getSessionInfo(
-                        rollback.stagedSessionId);
+                PackageInstaller.SessionInfo session =
+                        installer.getSessionInfo(rollback.getStagedSessionId());
                 if (session == null || session.isStagedSessionFailed()) {
                     // TODO: Do we need to remove this from
                     // mRollbacks, or is it okay to leave as
@@ -650,13 +650,13 @@
 
             for (Rollback rollback : restoreInProgress) {
                 PackageInstaller installer = mContext.getPackageManager().getPackageInstaller();
-                PackageInstaller.SessionInfo session = installer.getSessionInfo(
-                        rollback.stagedSessionId);
+                PackageInstaller.SessionInfo session =
+                        installer.getSessionInfo(rollback.getStagedSessionId());
                 // TODO: What if session is null?
                 if (session != null) {
                     if (session.isStagedSessionApplied() || session.isStagedSessionFailed()) {
                         synchronized (mLock) {
-                            rollback.restoreUserDataInProgress = false;
+                            rollback.setRestoreUserDataInProgress(false);
                         }
                         saveRollback(rollback);
                     }
@@ -694,8 +694,7 @@
             while (iter.hasNext()) {
                 Rollback rollback = iter.next();
                 // TODO: Should we remove rollbacks in the ENABLING state here?
-                if (rollback.state == Rollback.ROLLBACK_STATE_AVAILABLE
-                        || rollback.state == Rollback.ROLLBACK_STATE_ENABLING) {
+                if (rollback.isEnabling() || rollback.isAvailable()) {
                     for (PackageRollbackInfo info : rollback.info.getPackages()) {
                         if (info.getPackageName().equals(packageName)
                                 && !packageVersionsEqual(
@@ -761,15 +760,16 @@
             Iterator<Rollback> iter = mRollbacks.iterator();
             while (iter.hasNext()) {
                 Rollback rollback = iter.next();
-                if (rollback.state != Rollback.ROLLBACK_STATE_AVAILABLE) {
+                if (!rollback.isAvailable()) {
                     continue;
                 }
                 if (!now.isBefore(
-                            rollback.timestamp.plusMillis(mRollbackLifetimeDurationInMillis))) {
+                            rollback.getTimestamp()
+                                    .plusMillis(mRollbackLifetimeDurationInMillis))) {
                     iter.remove();
                     deleteRollback(rollback);
-                } else if (oldest == null || oldest.isAfter(rollback.timestamp)) {
-                    oldest = rollback.timestamp;
+                } else if (oldest == null || oldest.isAfter(rollback.getTimestamp())) {
+                    oldest = rollback.getTimestamp();
                 }
             }
         }
@@ -877,7 +877,7 @@
         synchronized (mLock) {
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 Rollback rollback = mRollbacks.get(i);
-                if (rollback.apkSessionId == parentSession.getSessionId()) {
+                if (rollback.getApkSessionId() == parentSession.getSessionId()) {
                     // This is the apk session for a staged session with rollback enabled. We do not
                     // need to create a new rollback for this session.
                     return true;
@@ -1020,7 +1020,7 @@
             // staged installs
             for (int i = 0; i < mRollbacks.size(); i++) {
                 Rollback rollback = mRollbacks.get(i);
-                if (rollback.state != Rollback.ROLLBACK_STATE_ENABLING) {
+                if (!rollback.isEnabling()) {
                     continue;
                 }
 
@@ -1053,7 +1053,7 @@
         synchronized (mLock) {
             for (int i = 0; i < mRollbacks.size(); ++i) {
                 Rollback candidate = mRollbacks.get(i);
-                if (candidate.restoreUserDataInProgress) {
+                if (candidate.isRestoreUserDataInProgress()) {
                     info = getPackageRollbackInfo(candidate, packageName);
                     if (info != null) {
                         rollback = candidate;
@@ -1146,8 +1146,8 @@
             synchronized (mLock) {
                 for (int i = 0; i < mRollbacks.size(); ++i) {
                     Rollback candidate = mRollbacks.get(i);
-                    if (candidate.stagedSessionId == originalSessionId) {
-                        candidate.apkSessionId = apkSessionId;
+                    if (candidate.getStagedSessionId() == originalSessionId) {
+                        candidate.setApkSessionId(apkSessionId);
                         rollback = candidate;
                         break;
                     }
@@ -1333,8 +1333,8 @@
         // to a new package being installed. Won't this revive an expired
         // rollback? Consider adding a ROLLBACK_STATE_EXPIRED to address this.
         synchronized (mLock) {
-            rollback.state = Rollback.ROLLBACK_STATE_AVAILABLE;
-            rollback.timestamp = Instant.now();
+            rollback.setAvailable();
+            rollback.setTimestamp(Instant.now());
         }
         saveRollback(rollback);
 
@@ -1434,9 +1434,9 @@
                 ipw.println(info.getRollbackId() + ":");
                 ipw.increaseIndent();
                 ipw.println("-state: " + rollback.getStateAsString());
-                ipw.println("-timestamp: " + rollback.timestamp);
-                if (rollback.stagedSessionId != -1) {
-                    ipw.println("-stagedSessionId: " + rollback.stagedSessionId);
+                ipw.println("-timestamp: " + rollback.getTimestamp());
+                if (rollback.getStagedSessionId() != -1) {
+                    ipw.println("-stagedSessionId: " + rollback.getStagedSessionId());
                 }
                 ipw.println("-packages:");
                 ipw.increaseIndent();
@@ -1446,7 +1446,7 @@
                             + " -> " + pkg.getVersionRolledBackTo().getLongVersionCode());
                 }
                 ipw.decreaseIndent();
-                if (rollback.state == Rollback.ROLLBACK_STATE_COMMITTED) {
+                if (rollback.isCommitted()) {
                     ipw.println("-causePackages:");
                     ipw.increaseIndent();
                     for (VersionedPackage cPkg : info.getCausePackages()) {
diff --git a/services/core/java/com/android/server/rollback/RollbackStore.java b/services/core/java/com/android/server/rollback/RollbackStore.java
index b2448f6..772c53f 100644
--- a/services/core/java/com/android/server/rollback/RollbackStore.java
+++ b/services/core/java/com/android/server/rollback/RollbackStore.java
@@ -17,7 +17,6 @@
 package com.android.server.rollback;
 
 import static com.android.server.rollback.Rollback.rollbackStateFromString;
-import static com.android.server.rollback.Rollback.rollbackStateToString;
 
 import android.annotation.NonNull;
 import android.content.pm.VersionedPackage;
@@ -216,7 +215,7 @@
     static void backupPackageCodePath(Rollback rollback, String packageName, String codePath)
             throws IOException {
         File sourceFile = new File(codePath);
-        File targetDir = new File(rollback.backupDir, packageName);
+        File targetDir = new File(rollback.getBackupDir(), packageName);
         targetDir.mkdirs();
         File targetFile = new File(targetDir, sourceFile.getName());
 
@@ -229,7 +228,7 @@
      * Includes the base apk and any splits. Returns null if none found.
      */
     static File[] getPackageCodePaths(Rollback rollback, String packageName) {
-        File targetDir = new File(rollback.backupDir, packageName);
+        File targetDir = new File(rollback.getBackupDir(), packageName);
         File[] files = targetDir.listFiles();
         if (files == null || files.length == 0) {
             return null;
@@ -243,7 +242,7 @@
      */
     static void deletePackageCodePaths(Rollback rollback) {
         for (PackageRollbackInfo info : rollback.info.getPackages()) {
-            File targetDir = new File(rollback.backupDir, info.getPackageName());
+            File targetDir = new File(rollback.getBackupDir(), info.getPackageName());
             removeFile(targetDir);
         }
     }
@@ -255,13 +254,13 @@
         try {
             JSONObject dataJson = new JSONObject();
             dataJson.put("info", rollbackInfoToJson(rollback.info));
-            dataJson.put("timestamp", rollback.timestamp.toString());
-            dataJson.put("stagedSessionId", rollback.stagedSessionId);
-            dataJson.put("state", rollbackStateToString(rollback.state));
-            dataJson.put("apkSessionId", rollback.apkSessionId);
-            dataJson.put("restoreUserDataInProgress", rollback.restoreUserDataInProgress);
+            dataJson.put("timestamp", rollback.getTimestamp().toString());
+            dataJson.put("stagedSessionId", rollback.getStagedSessionId());
+            dataJson.put("state", rollback.getStateAsString());
+            dataJson.put("apkSessionId", rollback.getApkSessionId());
+            dataJson.put("restoreUserDataInProgress", rollback.isRestoreUserDataInProgress());
 
-            PrintWriter pw = new PrintWriter(new File(rollback.backupDir, "rollback.json"));
+            PrintWriter pw = new PrintWriter(new File(rollback.getBackupDir(), "rollback.json"));
             pw.println(dataJson.toString());
             pw.close();
         } catch (JSONException e) {
@@ -273,7 +272,7 @@
      * Removes all persistent storage associated with the given rollback.
      */
     void deleteRollback(Rollback rollback) {
-        removeFile(rollback.backupDir);
+        removeFile(rollback.getBackupDir());
     }
 
     /**
diff --git a/services/tests/servicestests/src/com/android/server/rollback/RollbackUnitTest.java b/services/tests/servicestests/src/com/android/server/rollback/RollbackUnitTest.java
new file mode 100644
index 0000000..d27f1c7
--- /dev/null
+++ b/services/tests/servicestests/src/com/android/server/rollback/RollbackUnitTest.java
@@ -0,0 +1,77 @@
+/*
+ * Copyright (C) 2019 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.android.server.rollback;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.File;
+
+@RunWith(JUnit4.class)
+public class RollbackUnitTest {
+
+    @Test
+    public void newEmptyStagedRollbackDefaults() {
+        int rollbackId = 123;
+        int sessionId = 567;
+        File file = new File("/test/testing");
+
+        Rollback rollback = new Rollback(rollbackId, file, sessionId);
+
+        assertThat(rollback.isEnabling()).isTrue();
+        assertThat(rollback.getBackupDir().getAbsolutePath()).isEqualTo("/test/testing");
+        assertThat(rollback.isStaged()).isTrue();
+        assertThat(rollback.getStagedSessionId()).isEqualTo(567);
+    }
+
+    @Test
+    public void newEmptyNonStagedRollbackDefaults() {
+        int rollbackId = 123;
+        File file = new File("/test/testing");
+
+        Rollback rollback = new Rollback(rollbackId, file, -1);
+
+        assertThat(rollback.isEnabling()).isTrue();
+        assertThat(rollback.getBackupDir().getAbsolutePath()).isEqualTo("/test/testing");
+        assertThat(rollback.isStaged()).isFalse();
+    }
+
+    @Test
+    public void rollbackStateChanges() {
+        Rollback rollback = new Rollback(123, new File("/test/testing"), -1);
+
+        assertThat(rollback.isEnabling()).isTrue();
+        assertThat(rollback.isAvailable()).isFalse();
+        assertThat(rollback.isCommitted()).isFalse();
+
+        rollback.setAvailable();
+
+        assertThat(rollback.isEnabling()).isFalse();
+        assertThat(rollback.isAvailable()).isTrue();
+        assertThat(rollback.isCommitted()).isFalse();
+
+        rollback.setCommitted();
+
+        assertThat(rollback.isEnabling()).isFalse();
+        assertThat(rollback.isAvailable()).isFalse();
+        assertThat(rollback.isCommitted()).isTrue();
+    }
+
+}
diff --git a/startop/scripts/app_startup/app_startup_runner.py b/startop/scripts/app_startup/app_startup_runner.py
index eb582f9..fa1c4e6 100755
--- a/startop/scripts/app_startup/app_startup_runner.py
+++ b/startop/scripts/app_startup/app_startup_runner.py
@@ -41,11 +41,12 @@
 sys.path.append(os.path.dirname(DIR))
 import lib.cmd_utils as cmd_utils
 import lib.print_utils as print_utils
-import iorap.compiler as compiler
 from app_startup.run_app_with_prefetch import PrefetchAppRunner
 import app_startup.lib.args_utils as args_utils
 from app_startup.lib.data_frame import DataFrame
 from app_startup.lib.perfetto_trace_collector import PerfettoTraceCollector
+from iorap.compiler import CompilerType
+import iorap.compiler as compiler
 
 # The following command line options participate in the combinatorial generation.
 # All other arguments have a global effect.
@@ -58,8 +59,6 @@
 
 CollectorPackageInfo = NamedTuple('CollectorPackageInfo',
                                   [('package', str), ('compiler_filter', str)])
-_COMPILER_SCRIPT = os.path.join(os.path.dirname(os.path.dirname(
-    os.path.realpath(__file__))), 'iorap/compiler.py')
 # by 2; systrace starts up slowly.
 
 _UNLOCK_SCREEN_SCRIPT = os.path.join(
@@ -135,6 +134,10 @@
                               action='append',
                               help='The trace duration (milliseconds) in '
                                    'compilation')
+  optional_named.add_argument('--compiler-type', dest='compiler_type',
+                              type=CompilerType, choices=list(CompilerType),
+                              default=CompilerType.DEVICE,
+                              help='The type of compiler.')
 
   return parser.parse_args(argv)
 
@@ -211,26 +214,26 @@
 
   return DataFrame(d)
 
-def compile_perfetto_trace(inodes_path: str,
+def build_ri_compiler_argv(inodes_path: str,
                            perfetto_trace_file: str,
-                           trace_duration: Optional[timedelta]) -> TextIO:
-  compiler_trace_file = tempfile.NamedTemporaryFile()
-  argv = [_COMPILER_SCRIPT, '-i', inodes_path, '--perfetto-trace',
-          perfetto_trace_file, '-o', compiler_trace_file.name]
+                           trace_duration: Optional[timedelta]
+                           ) -> str:
+  argv = ['-i', inodes_path, '--perfetto-trace',
+          perfetto_trace_file]
 
   if trace_duration is not None:
     argv += ['--duration', str(int(trace_duration.total_seconds()
-                               * PerfettoTraceCollector.MS_PER_SEC))]
+                                   * PerfettoTraceCollector.MS_PER_SEC))]
 
   print_utils.debug_print(argv)
-  compiler.main(argv)
-  return compiler_trace_file
+  return argv
 
 def execute_run_using_perfetto_trace(collector_info,
                                      run_combos: Iterable[RunCommandArgs],
                                      simulate: bool,
                                      inodes_path: str,
-                                     timeout: int) -> DataFrame:
+                                     timeout: int,
+                                     compiler_type: CompilerType) -> DataFrame:
   """ Executes run based on perfetto trace. """
   passed, perfetto_trace_file = run_perfetto_collector(collector_info,
                                                        timeout,
@@ -244,9 +247,15 @@
         if simulate:
           compiler_trace_file = tempfile.NamedTemporaryFile()
         else:
-          compiler_trace_file = compile_perfetto_trace(inodes_path,
-                                                       perfetto_trace_file.name,
-                                                       combos.trace_duration)
+          ri_compiler_argv = build_ri_compiler_argv(inodes_path,
+                                                    perfetto_trace_file.name,
+                                                    combos.trace_duration)
+          compiler_trace_file = compiler.compile(compiler_type,
+                                                 inodes_path,
+                                                 ri_compiler_argv,
+                                                 combos.package,
+                                                 combos.activity)
+
         with compiler_trace_file:
           combos = combos._replace(input=compiler_trace_file.name)
           print_utils.debug_print(combos)
@@ -261,7 +270,8 @@
     grouped_run_combos: Iterable[Tuple[CollectorPackageInfo, Iterable[RunCommandArgs]]],
     simulate: bool,
     inodes_path: str,
-    timeout: int):
+    timeout: int,
+    compiler_type: CompilerType):
   # nothing will work if the screen isn't unlocked first.
   cmd_utils.execute_arbitrary_command([_UNLOCK_SCREEN_SCRIPT],
                                       timeout,
@@ -273,7 +283,8 @@
                                                 run_combos,
                                                 simulate,
                                                 inodes_path,
-                                                timeout)
+                                                timeout,
+                                                compiler_type)
 
 def gather_results(commands: Iterable[Tuple[DataFrame]],
                    key_list: List[str], value_list: List[Tuple[str, ...]]):
@@ -361,7 +372,8 @@
   exec = execute_run_combos(grouped_combos(),
                             opts.simulate,
                             opts.inodes,
-                            opts.timeout)
+                            opts.timeout,
+                            opts.compiler_type)
 
   results = gather_results(exec, _COMBINATORIAL_OPTIONS, combos())
 
diff --git a/startop/scripts/app_startup/app_startup_runner_test.py b/startop/scripts/app_startup/app_startup_runner_test.py
index 42ea5f0..382f6f3 100755
--- a/startop/scripts/app_startup/app_startup_runner_test.py
+++ b/startop/scripts/app_startup/app_startup_runner_test.py
@@ -92,7 +92,7 @@
   """
   d = {'compiler_filters': None, 'simulate': False, 'debug': False,
        'output': None, 'timeout': 10, 'loop_count': 1, 'inodes': None,
-       'trace_duration': None}
+       'trace_duration': None, 'compiler_type': asr.CompilerType.HOST}
   d.update(kwargs)
   return d
 
diff --git a/startop/scripts/iorap/compiler.py b/startop/scripts/iorap/compiler.py
old mode 100755
new mode 100644
index 17b58c1..1426d34
--- a/startop/scripts/iorap/compiler.py
+++ b/startop/scripts/iorap/compiler.py
@@ -1,323 +1,73 @@
 #!/usr/bin/env python3
-
 #
-# Copyright (C) 2019 The Android Open Source Project
+# Copyright 2019, 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
+#     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.
-#
 
-#
-# Dependencies:
-#
-# $> sudo apt-get install python3-pip
-# $> pip3 install --user protobuf sqlalchemy sqlite3
-#
-
-import optparse
+import importlib
 import os
-import re
 import sys
 import tempfile
-from pathlib import Path
-from datetime import timedelta
-from typing import Iterable, Optional, List
+from enum import Enum
+from typing import TextIO, List
 
+# local import
 DIR = os.path.abspath(os.path.dirname(__file__))
 sys.path.append(os.path.dirname(DIR))
-from iorap.generated.TraceFile_pb2 import *
-from iorap.lib.inode2filename import Inode2Filename
+import lib.print_utils as print_utils
 
-parent_dir_name = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
-sys.path.append(parent_dir_name)
-from trace_analyzer.lib.trace2db import Trace2Db, MmFilemapAddToPageCache, \
-    RawFtraceEntry
-import lib.cmd_utils as cmd_utils
+# Type of compiler.
+class CompilerType(Enum):
+  HOST = 1  # iorap.cmd.compiler on host
+  DEVICE = 2  # adb shell iorap.cmd.compiler
+  RI = 3  # compiler.py
 
-_PAGE_SIZE = 4096 # adb shell getconf PAGESIZE ## size of a memory page in bytes.
-ANDROID_BUILD_TOP = Path(parent_dir_name).parents[3]
-TRACECONV_BIN = ANDROID_BUILD_TOP.joinpath(
-    'external/perfetto/tools/traceconv')
+def compile_perfetto_trace_ri(
+    argv: List[str],
+    compiler) -> TextIO:
+  print_utils.debug_print('Compile using RI compiler.')
+  compiler_trace_file = tempfile.NamedTemporaryFile()
+  argv.extend(['-o', compiler_trace_file.name])
+  print_utils.debug_print(argv)
+  compiler.main([''] + argv)
+  return compiler_trace_file
 
-class PageRun:
-  """
-  Intermediate representation for a run of one or more pages.
-  """
-  def __init__(self, device_number: int, inode: int, offset: int, length: int):
-    self.device_number = device_number
-    self.inode = inode
-    self.offset = offset
-    self.length = length
+def compile_perfetto_trace_device(inodes_path: str,
+                                  package: str,
+                                  activity: str,
+                                  compiler) -> TextIO:
+  print_utils.debug_print('Compile using on-device compiler.')
+  compiler_trace_file = tempfile.NamedTemporaryFile()
+  compiler.main(inodes_path, package, activity, compiler_trace_file.name)
+  return compiler_trace_file
 
-  def __str__(self):
-    return "PageRun(device_number=%d, inode=%d, offset=%d, length=%d)" \
-        %(self.device_number, self.inode, self.offset, self.length)
+def compile(compiler_type: CompilerType,
+            inodes_path: str,
+            ri_compiler_argv,
+            package: str,
+            activity: str) -> TextIO:
+  if compiler_type == CompilerType.RI:
+    compiler = importlib.import_module('iorap.compiler_ri')
+    compiler_trace_file = compile_perfetto_trace_ri(ri_compiler_argv,
+                                                    compiler)
+    return compiler_trace_file
+  if compiler_type == CompilerType.DEVICE:
+    compiler = importlib.import_module('iorap.compiler_device')
+    compiler_trace_file = compile_perfetto_trace_device(inodes_path,
+                                                        package,
+                                                        activity,
+                                                        compiler)
+    return compiler_trace_file
 
-def debug_print(msg):
-  #print(msg)
-  pass
-
-UNDER_LAUNCH = False
-
-def page_cache_entries_to_runs(page_cache_entries: Iterable[MmFilemapAddToPageCache]):
-  global _PAGE_SIZE
-
-  runs = [
-      PageRun(device_number=pg_entry.dev, inode=pg_entry.ino, offset=pg_entry.ofs,
-              length=_PAGE_SIZE)
-        for pg_entry in page_cache_entries
-  ]
-
-  for r in runs:
-    debug_print(r)
-
-  print("Stats: Page runs totaling byte length: %d" %(len(runs) * _PAGE_SIZE))
-
-  return runs
-
-def optimize_page_runs(page_runs):
-  new_entries = []
-  last_entry = None
-  for pg_entry in page_runs:
-    if last_entry:
-      if pg_entry.device_number == last_entry.device_number and pg_entry.inode == last_entry.inode:
-        # we are dealing with a run for the same exact file as a previous run.
-        if pg_entry.offset == last_entry.offset + last_entry.length:
-          # trivially contiguous entries. merge them together.
-          last_entry.length += pg_entry.length
-          continue
-    # Default: Add the run without merging it to a previous run.
-    last_entry = pg_entry
-    new_entries.append(pg_entry)
-  return new_entries
-
-def is_filename_matching_filter(file_name, filters=[]):
-  """
-  Blacklist-style regular expression filters.
-
-  :return: True iff file_name has an RE match in one of the filters.
-  """
-  for filt in filters:
-    res = re.search(filt, file_name)
-    if res:
-      return True
-
-  return False
-
-def build_protobuf(page_runs, inode2filename, filters=[]):
-  trace_file = TraceFile()
-  trace_file_index = trace_file.index
-
-  file_id_counter = 0
-  file_id_map = {} # filename -> id
-
-  stats_length_total = 0
-  filename_stats = {} # filename -> total size
-
-  skipped_inode_map = {}
-  filtered_entry_map = {} # filename -> count
-
-  for pg_entry in page_runs:
-    fn = inode2filename.resolve(pg_entry.device_number, pg_entry.inode)
-    if not fn:
-      skipped_inode_map[pg_entry.inode] = skipped_inode_map.get(pg_entry.inode, 0) + 1
-      continue
-
-    filename = fn
-
-    if filters and not is_filename_matching_filter(filename, filters):
-      filtered_entry_map[filename] = filtered_entry_map.get(filename, 0) + 1
-      continue
-
-    file_id = file_id_map.get(filename)
-    if not file_id:
-      file_id = file_id_counter
-      file_id_map[filename] = file_id_counter
-      file_id_counter = file_id_counter + 1
-
-      file_index_entry = trace_file_index.entries.add()
-      file_index_entry.id = file_id
-      file_index_entry.file_name = filename
-
-    # already in the file index, add the file entry.
-    file_entry = trace_file.list.entries.add()
-    file_entry.index_id = file_id
-    file_entry.file_length = pg_entry.length
-    stats_length_total += file_entry.file_length
-    file_entry.file_offset = pg_entry.offset
-
-    filename_stats[filename] = filename_stats.get(filename, 0) + file_entry.file_length
-
-  for inode, count in skipped_inode_map.items():
-    print("WARNING: Skip inode %s because it's not in inode map (%d entries)" %(inode, count))
-
-  print("Stats: Sum of lengths %d" %(stats_length_total))
-
-  if filters:
-    print("Filter: %d total files removed." %(len(filtered_entry_map)))
-
-    for fn, count in filtered_entry_map.items():
-      print("Filter: File '%s' removed '%d' entries." %(fn, count))
-
-  for filename, file_size in filename_stats.items():
-    print("%s,%s" %(filename, file_size))
-
-  return trace_file
-
-def calc_trace_end_time(trace2db: Trace2Db,
-                        trace_duration: Optional[timedelta]) -> float:
-  """
-  Calculates the end time based on the trace duration.
-  The start time is the first receiving mm file map event.
-  The end time is the start time plus the trace duration.
-  All of them are in milliseconds.
-  """
-  # If the duration is not set, assume all time is acceptable.
-  if trace_duration is None:
-    # float('inf')
-    return RawFtraceEntry.__table__.c.timestamp.type.python_type('inf')
-
-  first_event = trace2db.session.query(MmFilemapAddToPageCache).join(
-      MmFilemapAddToPageCache.raw_ftrace_entry).order_by(
-      RawFtraceEntry.timestamp).first()
-
-  # total_seconds() will return a float number.
-  return first_event.raw_ftrace_entry.timestamp + trace_duration.total_seconds()
-
-def query_add_to_page_cache(trace2db: Trace2Db, trace_duration: Optional[timedelta]):
-  end_time = calc_trace_end_time(trace2db, trace_duration)
-  # SELECT * FROM tbl ORDER BY id;
-  return trace2db.session.query(MmFilemapAddToPageCache).join(
-      MmFilemapAddToPageCache.raw_ftrace_entry).filter(
-      RawFtraceEntry.timestamp <= end_time).order_by(
-      MmFilemapAddToPageCache.id).all()
-
-def transform_perfetto_trace_to_systrace(path_to_perfetto_trace: str,
-                                         path_to_tmp_systrace: str) -> None:
-  """ Transforms the systrace file from perfetto trace. """
-  cmd_utils.run_command_nofail([str(TRACECONV_BIN),
-                                'systrace',
-                                path_to_perfetto_trace,
-                                path_to_tmp_systrace])
-
-
-def run(sql_db_path:str,
-        trace_file:str,
-        trace_duration:Optional[timedelta],
-        output_file:str,
-        inode_table:str,
-        filter:List[str]) -> int:
-  trace2db = Trace2Db(sql_db_path)
-  # Speed optimization: Skip any entries that aren't mm_filemap_add_to_pagecache.
-  trace2db.set_raw_ftrace_entry_filter(\
-      lambda entry: entry['function'] == 'mm_filemap_add_to_page_cache')
-  # TODO: parse multiple trace files here.
-  parse_count = trace2db.parse_file_into_db(trace_file)
-
-  mm_filemap_add_to_page_cache_rows = query_add_to_page_cache(trace2db,
-                                                              trace_duration)
-  print("DONE. Parsed %d entries into sql db." %(len(mm_filemap_add_to_page_cache_rows)))
-
-  page_runs = page_cache_entries_to_runs(mm_filemap_add_to_page_cache_rows)
-  print("DONE. Converted %d entries" %(len(page_runs)))
-
-  # TODO: flags to select optimizations.
-  optimized_page_runs = optimize_page_runs(page_runs)
-  print("DONE. Optimized down to %d entries" %(len(optimized_page_runs)))
-
-  print("Build protobuf...")
-  trace_file = build_protobuf(optimized_page_runs, inode_table, filter)
-
-  print("Write protobuf to file...")
-  output_file = open(output_file, 'wb')
-  output_file.write(trace_file.SerializeToString())
-  output_file.close()
-
-  print("DONE")
-
-  # TODO: Silent running mode [no output except on error] for build runs.
-
-  return 0
-
-def main(argv):
-  parser = optparse.OptionParser(usage="Usage: %prog [options]", description="Compile systrace file into TraceFile.pb")
-  parser.add_option('-i', dest='inode_data_file', metavar='FILE',
-                    help='Read cached inode data from a file saved earlier with pagecache.py -d')
-  parser.add_option('-t', dest='trace_file', metavar='FILE',
-                    help='Path to systrace file (trace.html) that will be parsed')
-  parser.add_option('--perfetto-trace', dest='perfetto_trace_file',
-                    metavar='FILE',
-                    help='Path to perfetto trace that will be parsed')
-
-  parser.add_option('--db', dest='sql_db', metavar='FILE',
-                    help='Path to intermediate sqlite3 database [default: in-memory].')
-
-  parser.add_option('-f', dest='filter', action="append", default=[],
-                    help="Add file filter. All file entries not matching one of the filters are discarded.")
-
-  parser.add_option('-l', dest='launch_lock', action="store_true", default=False,
-                    help="Exclude all events not inside launch_lock")
-
-  parser.add_option('-o', dest='output_file', metavar='FILE',
-                    help='Output protobuf file')
-
-  parser.add_option('--duration', dest='trace_duration', action="store",
-                    type=int, help='The duration of trace in milliseconds.')
-
-  options, categories = parser.parse_args(argv[1:])
-
-  # TODO: OptionParser should have some flags to make these mandatory.
-  if not options.inode_data_file:
-    parser.error("-i is required")
-  if not options.trace_file and not options.perfetto_trace_file:
-    parser.error("one of -t or --perfetto-trace is required")
-  if options.trace_file and options.perfetto_trace_file:
-    parser.error("please enter either -t or --perfetto-trace, not both")
-  if not options.output_file:
-    parser.error("-o is required")
-
-  if options.launch_lock:
-    print("INFO: Launch lock flag (-l) enabled; filtering all events not inside launch_lock.")
-
-  inode_table = Inode2Filename.new_from_filename(options.inode_data_file)
-
-  sql_db_path = ":memory:"
-  if options.sql_db:
-    sql_db_path = options.sql_db
-
-  trace_duration = timedelta(milliseconds=options.trace_duration) if \
-    options.trace_duration is not None else None
-
-  # if the input is systrace
-  if options.trace_file:
-    return run(sql_db_path,
-               options.trace_file,
-               trace_duration,
-               options.output_file,
-               inode_table,
-               options.filter)
-
-  # if the input is perfetto trace
-  # TODO python 3.7 switch to using nullcontext
-  with tempfile.NamedTemporaryFile() as trace_file:
-    transform_perfetto_trace_to_systrace(options.perfetto_trace_file,
-                                         trace_file.name)
-    return run(sql_db_path,
-               trace_file.name,
-               trace_duration,
-               options.output_file,
-               inode_table,
-               options.filter)
-
-if __name__ == '__main__':
-  print(sys.argv)
-  sys.exit(main(sys.argv))
+  # Should not arrive here.
+  raise ValueError('Unknown compiler type')
diff --git a/startop/scripts/iorap/compiler_device.py b/startop/scripts/iorap/compiler_device.py
new file mode 100644
index 0000000..d941cd9
--- /dev/null
+++ b/startop/scripts/iorap/compiler_device.py
@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+#
+# Copyright 2019, 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.
+
+import argparse
+import os
+import sys
+from typing import List
+
+DIR = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.dirname(DIR))  # framework/base/startop/script
+import lib.print_utils as print_utils
+import iorap.lib.iorapd_utils as iorapd_utils
+from app_startup.lib.app_runner import AppRunner
+
+IORAP_COMMON_BASH_SCRIPT = os.path.join(DIR, 'common')
+
+def parse_options(argv: List[str] = None):
+  """Parses command line arguments and returns an argparse Namespace object."""
+  parser = argparse.ArgumentParser(description="Compile perfetto trace file")
+  required_named = parser.add_argument_group('required named arguments')
+
+  required_named.add_argument('-i', dest='inodes', metavar='FILE',
+                              help='Read cached inode data from a file saved '
+                                   'earlier with pagecache.py -d')
+  required_named.add_argument('-p', dest='package',
+                              help='Package of the app to be compiled')
+
+  optional_named = parser.add_argument_group('optional named arguments')
+  optional_named.add_argument('-o', dest='output',
+                              help='The compiled trace is stored into the output file')
+  optional_named.add_argument('-a', dest='activity',
+                              help='Activity of the app to be compiled')
+  optional_named.add_argument('-d', dest='debug', action='store_true'
+                              , help='Activity of the app to be compiled')
+
+  return parser.parse_args(argv)
+
+def main(inodes, package, activity, output, **kwargs) -> int:
+  """Entries of the program."""
+  if not activity:
+    activity = AppRunner.get_activity(package)
+
+  passed = iorapd_utils.compile_perfetto_trace_on_device(package, activity,
+                                                         inodes)
+  if passed and output:
+    iorapd_utils.get_iorapd_compiler_trace(package, activity, output)
+
+  return 0
+
+if __name__ == '__main__':
+  opts = parse_options()
+  if opts.debug:
+    print_utils.DEBUG = opts.debug
+  print_utils.debug_print(opts)
+  sys.exit(main(**(vars(opts))))
diff --git a/startop/scripts/iorap/compiler_ri.py b/startop/scripts/iorap/compiler_ri.py
new file mode 100755
index 0000000..17b58c1
--- /dev/null
+++ b/startop/scripts/iorap/compiler_ri.py
@@ -0,0 +1,323 @@
+#!/usr/bin/env python3
+
+#
+# Copyright (C) 2019 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.
+#
+
+#
+# Dependencies:
+#
+# $> sudo apt-get install python3-pip
+# $> pip3 install --user protobuf sqlalchemy sqlite3
+#
+
+import optparse
+import os
+import re
+import sys
+import tempfile
+from pathlib import Path
+from datetime import timedelta
+from typing import Iterable, Optional, List
+
+DIR = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.dirname(DIR))
+from iorap.generated.TraceFile_pb2 import *
+from iorap.lib.inode2filename import Inode2Filename
+
+parent_dir_name = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
+sys.path.append(parent_dir_name)
+from trace_analyzer.lib.trace2db import Trace2Db, MmFilemapAddToPageCache, \
+    RawFtraceEntry
+import lib.cmd_utils as cmd_utils
+
+_PAGE_SIZE = 4096 # adb shell getconf PAGESIZE ## size of a memory page in bytes.
+ANDROID_BUILD_TOP = Path(parent_dir_name).parents[3]
+TRACECONV_BIN = ANDROID_BUILD_TOP.joinpath(
+    'external/perfetto/tools/traceconv')
+
+class PageRun:
+  """
+  Intermediate representation for a run of one or more pages.
+  """
+  def __init__(self, device_number: int, inode: int, offset: int, length: int):
+    self.device_number = device_number
+    self.inode = inode
+    self.offset = offset
+    self.length = length
+
+  def __str__(self):
+    return "PageRun(device_number=%d, inode=%d, offset=%d, length=%d)" \
+        %(self.device_number, self.inode, self.offset, self.length)
+
+def debug_print(msg):
+  #print(msg)
+  pass
+
+UNDER_LAUNCH = False
+
+def page_cache_entries_to_runs(page_cache_entries: Iterable[MmFilemapAddToPageCache]):
+  global _PAGE_SIZE
+
+  runs = [
+      PageRun(device_number=pg_entry.dev, inode=pg_entry.ino, offset=pg_entry.ofs,
+              length=_PAGE_SIZE)
+        for pg_entry in page_cache_entries
+  ]
+
+  for r in runs:
+    debug_print(r)
+
+  print("Stats: Page runs totaling byte length: %d" %(len(runs) * _PAGE_SIZE))
+
+  return runs
+
+def optimize_page_runs(page_runs):
+  new_entries = []
+  last_entry = None
+  for pg_entry in page_runs:
+    if last_entry:
+      if pg_entry.device_number == last_entry.device_number and pg_entry.inode == last_entry.inode:
+        # we are dealing with a run for the same exact file as a previous run.
+        if pg_entry.offset == last_entry.offset + last_entry.length:
+          # trivially contiguous entries. merge them together.
+          last_entry.length += pg_entry.length
+          continue
+    # Default: Add the run without merging it to a previous run.
+    last_entry = pg_entry
+    new_entries.append(pg_entry)
+  return new_entries
+
+def is_filename_matching_filter(file_name, filters=[]):
+  """
+  Blacklist-style regular expression filters.
+
+  :return: True iff file_name has an RE match in one of the filters.
+  """
+  for filt in filters:
+    res = re.search(filt, file_name)
+    if res:
+      return True
+
+  return False
+
+def build_protobuf(page_runs, inode2filename, filters=[]):
+  trace_file = TraceFile()
+  trace_file_index = trace_file.index
+
+  file_id_counter = 0
+  file_id_map = {} # filename -> id
+
+  stats_length_total = 0
+  filename_stats = {} # filename -> total size
+
+  skipped_inode_map = {}
+  filtered_entry_map = {} # filename -> count
+
+  for pg_entry in page_runs:
+    fn = inode2filename.resolve(pg_entry.device_number, pg_entry.inode)
+    if not fn:
+      skipped_inode_map[pg_entry.inode] = skipped_inode_map.get(pg_entry.inode, 0) + 1
+      continue
+
+    filename = fn
+
+    if filters and not is_filename_matching_filter(filename, filters):
+      filtered_entry_map[filename] = filtered_entry_map.get(filename, 0) + 1
+      continue
+
+    file_id = file_id_map.get(filename)
+    if not file_id:
+      file_id = file_id_counter
+      file_id_map[filename] = file_id_counter
+      file_id_counter = file_id_counter + 1
+
+      file_index_entry = trace_file_index.entries.add()
+      file_index_entry.id = file_id
+      file_index_entry.file_name = filename
+
+    # already in the file index, add the file entry.
+    file_entry = trace_file.list.entries.add()
+    file_entry.index_id = file_id
+    file_entry.file_length = pg_entry.length
+    stats_length_total += file_entry.file_length
+    file_entry.file_offset = pg_entry.offset
+
+    filename_stats[filename] = filename_stats.get(filename, 0) + file_entry.file_length
+
+  for inode, count in skipped_inode_map.items():
+    print("WARNING: Skip inode %s because it's not in inode map (%d entries)" %(inode, count))
+
+  print("Stats: Sum of lengths %d" %(stats_length_total))
+
+  if filters:
+    print("Filter: %d total files removed." %(len(filtered_entry_map)))
+
+    for fn, count in filtered_entry_map.items():
+      print("Filter: File '%s' removed '%d' entries." %(fn, count))
+
+  for filename, file_size in filename_stats.items():
+    print("%s,%s" %(filename, file_size))
+
+  return trace_file
+
+def calc_trace_end_time(trace2db: Trace2Db,
+                        trace_duration: Optional[timedelta]) -> float:
+  """
+  Calculates the end time based on the trace duration.
+  The start time is the first receiving mm file map event.
+  The end time is the start time plus the trace duration.
+  All of them are in milliseconds.
+  """
+  # If the duration is not set, assume all time is acceptable.
+  if trace_duration is None:
+    # float('inf')
+    return RawFtraceEntry.__table__.c.timestamp.type.python_type('inf')
+
+  first_event = trace2db.session.query(MmFilemapAddToPageCache).join(
+      MmFilemapAddToPageCache.raw_ftrace_entry).order_by(
+      RawFtraceEntry.timestamp).first()
+
+  # total_seconds() will return a float number.
+  return first_event.raw_ftrace_entry.timestamp + trace_duration.total_seconds()
+
+def query_add_to_page_cache(trace2db: Trace2Db, trace_duration: Optional[timedelta]):
+  end_time = calc_trace_end_time(trace2db, trace_duration)
+  # SELECT * FROM tbl ORDER BY id;
+  return trace2db.session.query(MmFilemapAddToPageCache).join(
+      MmFilemapAddToPageCache.raw_ftrace_entry).filter(
+      RawFtraceEntry.timestamp <= end_time).order_by(
+      MmFilemapAddToPageCache.id).all()
+
+def transform_perfetto_trace_to_systrace(path_to_perfetto_trace: str,
+                                         path_to_tmp_systrace: str) -> None:
+  """ Transforms the systrace file from perfetto trace. """
+  cmd_utils.run_command_nofail([str(TRACECONV_BIN),
+                                'systrace',
+                                path_to_perfetto_trace,
+                                path_to_tmp_systrace])
+
+
+def run(sql_db_path:str,
+        trace_file:str,
+        trace_duration:Optional[timedelta],
+        output_file:str,
+        inode_table:str,
+        filter:List[str]) -> int:
+  trace2db = Trace2Db(sql_db_path)
+  # Speed optimization: Skip any entries that aren't mm_filemap_add_to_pagecache.
+  trace2db.set_raw_ftrace_entry_filter(\
+      lambda entry: entry['function'] == 'mm_filemap_add_to_page_cache')
+  # TODO: parse multiple trace files here.
+  parse_count = trace2db.parse_file_into_db(trace_file)
+
+  mm_filemap_add_to_page_cache_rows = query_add_to_page_cache(trace2db,
+                                                              trace_duration)
+  print("DONE. Parsed %d entries into sql db." %(len(mm_filemap_add_to_page_cache_rows)))
+
+  page_runs = page_cache_entries_to_runs(mm_filemap_add_to_page_cache_rows)
+  print("DONE. Converted %d entries" %(len(page_runs)))
+
+  # TODO: flags to select optimizations.
+  optimized_page_runs = optimize_page_runs(page_runs)
+  print("DONE. Optimized down to %d entries" %(len(optimized_page_runs)))
+
+  print("Build protobuf...")
+  trace_file = build_protobuf(optimized_page_runs, inode_table, filter)
+
+  print("Write protobuf to file...")
+  output_file = open(output_file, 'wb')
+  output_file.write(trace_file.SerializeToString())
+  output_file.close()
+
+  print("DONE")
+
+  # TODO: Silent running mode [no output except on error] for build runs.
+
+  return 0
+
+def main(argv):
+  parser = optparse.OptionParser(usage="Usage: %prog [options]", description="Compile systrace file into TraceFile.pb")
+  parser.add_option('-i', dest='inode_data_file', metavar='FILE',
+                    help='Read cached inode data from a file saved earlier with pagecache.py -d')
+  parser.add_option('-t', dest='trace_file', metavar='FILE',
+                    help='Path to systrace file (trace.html) that will be parsed')
+  parser.add_option('--perfetto-trace', dest='perfetto_trace_file',
+                    metavar='FILE',
+                    help='Path to perfetto trace that will be parsed')
+
+  parser.add_option('--db', dest='sql_db', metavar='FILE',
+                    help='Path to intermediate sqlite3 database [default: in-memory].')
+
+  parser.add_option('-f', dest='filter', action="append", default=[],
+                    help="Add file filter. All file entries not matching one of the filters are discarded.")
+
+  parser.add_option('-l', dest='launch_lock', action="store_true", default=False,
+                    help="Exclude all events not inside launch_lock")
+
+  parser.add_option('-o', dest='output_file', metavar='FILE',
+                    help='Output protobuf file')
+
+  parser.add_option('--duration', dest='trace_duration', action="store",
+                    type=int, help='The duration of trace in milliseconds.')
+
+  options, categories = parser.parse_args(argv[1:])
+
+  # TODO: OptionParser should have some flags to make these mandatory.
+  if not options.inode_data_file:
+    parser.error("-i is required")
+  if not options.trace_file and not options.perfetto_trace_file:
+    parser.error("one of -t or --perfetto-trace is required")
+  if options.trace_file and options.perfetto_trace_file:
+    parser.error("please enter either -t or --perfetto-trace, not both")
+  if not options.output_file:
+    parser.error("-o is required")
+
+  if options.launch_lock:
+    print("INFO: Launch lock flag (-l) enabled; filtering all events not inside launch_lock.")
+
+  inode_table = Inode2Filename.new_from_filename(options.inode_data_file)
+
+  sql_db_path = ":memory:"
+  if options.sql_db:
+    sql_db_path = options.sql_db
+
+  trace_duration = timedelta(milliseconds=options.trace_duration) if \
+    options.trace_duration is not None else None
+
+  # if the input is systrace
+  if options.trace_file:
+    return run(sql_db_path,
+               options.trace_file,
+               trace_duration,
+               options.output_file,
+               inode_table,
+               options.filter)
+
+  # if the input is perfetto trace
+  # TODO python 3.7 switch to using nullcontext
+  with tempfile.NamedTemporaryFile() as trace_file:
+    transform_perfetto_trace_to_systrace(options.perfetto_trace_file,
+                                         trace_file.name)
+    return run(sql_db_path,
+               trace_file.name,
+               trace_duration,
+               options.output_file,
+               inode_table,
+               options.filter)
+
+if __name__ == '__main__':
+  print(sys.argv)
+  sys.exit(main(sys.argv))
diff --git a/startop/scripts/iorap/compiler_test.py b/startop/scripts/iorap/compiler_test.py
index 1a9f059..d1f11c5 100644
--- a/startop/scripts/iorap/compiler_test.py
+++ b/startop/scripts/iorap/compiler_test.py
@@ -30,7 +30,7 @@
 """
 import os
 
-import compiler
+import compiler_host as compiler
 
 DIR = os.path.abspath(os.path.dirname(__file__))
 TEXTCACHE = os.path.join(DIR, 'test_fixtures/compiler/common_textcache')
diff --git a/startop/scripts/iorap/lib/iorapd_utils.py b/startop/scripts/iorap/lib/iorapd_utils.py
index 0d62180..f6f21fd 100644
--- a/startop/scripts/iorap/lib/iorapd_utils.py
+++ b/startop/scripts/iorap/lib/iorapd_utils.py
@@ -18,10 +18,9 @@
 
 import os
 import sys
-from pathlib import Path
 
-# up to two level, like '../../'
-sys.path.append(Path(os.path.abspath(__file__)).parents[2])
+# up to two level
+sys.path.append(os.path.join(os.path.abspath(__file__),'../..'))
 import lib.cmd_utils as cmd_utils
 
 IORAPID_LIB_DIR = os.path.abspath(os.path.dirname(__file__))
@@ -39,6 +38,22 @@
   # Match logic of 'AppComponentName' in iorap::compiler C++ code.
   return '{}/{}%2F{}.{}'.format(IORAPD_DATA_PATH, package, activity, suffix)
 
+def compile_perfetto_trace_on_device(package: str, activity: str,
+                                     inodes: str) -> bool:
+  """Compiles the perfetto trace using on-device compiler."""
+  passed, _ = cmd_utils.run_shell_func(IORAP_COMMON_BASH_SCRIPT,
+                                       'iorapd_compiler_for_app_trace',
+                                       [package, activity, inodes])
+  return passed
+
+def get_iorapd_compiler_trace(package: str, activity: str, dest: str) -> str:
+  """Gets compiler trace to dest file."""
+  src = _iorapd_path_to_data_file(package, activity, 'compiled_trace.pb')
+  passed, _ = cmd_utils.run_shell_command('adb pull "{}" "{}"'.format(src, dest))
+  if not passed:
+    return False
+  return True
+
 def iorapd_compiler_install_trace_file(package: str, activity: str,
                                        input_file: str) -> bool:
   """Installs a compiled trace file.
diff --git a/tests/BootImageProfileTest/Android.bp b/tests/BootImageProfileTest/Android.bp
new file mode 100644
index 0000000..1b097a8
--- /dev/null
+++ b/tests/BootImageProfileTest/Android.bp
@@ -0,0 +1,20 @@
+// Copyright (C) 2019 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.
+
+java_test_host {
+    name: "BootImageProfileTest",
+    srcs: ["src/**/*.java"],
+    libs: ["tradefed"],
+    test_suites: ["general-tests"],
+}
diff --git a/tests/BootImageProfileTest/AndroidTest.xml b/tests/BootImageProfileTest/AndroidTest.xml
new file mode 100644
index 0000000..c132007
--- /dev/null
+++ b/tests/BootImageProfileTest/AndroidTest.xml
@@ -0,0 +1,39 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2019 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.
+-->
+<configuration description="Config for BootImageProfileTest">
+     <!-- do not use DeviceSetup#set-property because it reboots the device b/136200738.
+         furthermore the changes in /data/local.prop don't actually seem to get picked up.
+    -->
+    <target_preparer
+        class="com.android.tradefed.targetprep.DeviceSetup">
+        <!-- we need this magic flag, otherwise it always reboots and breaks the selinux -->
+        <option name="force-skip-system-props" value="true" />
+
+        <option name="run-command" value="setprop dalvik.vm.profilesystemserver true" />
+        <option name="run-command" value="setprop dalvik.vm.profilebootclasspath true" />
+
+        <!-- Profiling does not pick up the above changes we restart the shell -->
+        <option name="run-command" value="stop" />
+        <option name="run-command" value="start" />
+
+        <!-- give it some time to restart the shell; otherwise the first unit test might fail -->
+        <option name="run-command" value="sleep 2" />
+    </target_preparer>
+
+    <test class="com.android.tradefed.testtype.HostTest" >
+        <option name="class" value="com.android.bootimageprofile.BootImageProfileTest" />
+    </test>
+</configuration>
diff --git a/tests/BootImageProfileTest/TEST_MAPPING b/tests/BootImageProfileTest/TEST_MAPPING
new file mode 100644
index 0000000..1b569f9
--- /dev/null
+++ b/tests/BootImageProfileTest/TEST_MAPPING
@@ -0,0 +1,7 @@
+{
+  "presubmit": [
+    {
+      "name": "BootImageProfileTest"
+    }
+  ]
+}
diff --git a/tests/BootImageProfileTest/src/com/android/bootimageprofile/BootImageProfileTest.java b/tests/BootImageProfileTest/src/com/android/bootimageprofile/BootImageProfileTest.java
new file mode 100644
index 0000000..17986a3
--- /dev/null
+++ b/tests/BootImageProfileTest/src/com/android/bootimageprofile/BootImageProfileTest.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright (C) 2019 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.android.bootimageprofile;
+
+import static org.junit.Assert.assertTrue;
+
+import com.android.tradefed.device.ITestDevice;
+import com.android.tradefed.testtype.DeviceJUnit4ClassRunner;
+import com.android.tradefed.testtype.IDeviceTest;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(DeviceJUnit4ClassRunner.class)
+public class BootImageProfileTest implements IDeviceTest {
+    private ITestDevice mTestDevice;
+    private static final String SYSTEM_SERVER_PROFILE =
+            "/data/misc/profiles/cur/0/android/primary.prof";
+
+    @Override
+    public void setDevice(ITestDevice testDevice) {
+        mTestDevice = testDevice;
+    }
+
+    @Override
+    public ITestDevice getDevice() {
+        return mTestDevice;
+    }
+
+    /**
+     * Test that the boot image profile properties are set.
+     */
+    @Test
+    public void testProperties() throws Exception {
+        String res = mTestDevice.getProperty("dalvik.vm.profilebootclasspath");
+        assertTrue("profile boot class path not enabled", res != null && res.equals("true"));
+        res = mTestDevice.getProperty("dalvik.vm.profilesystemserver");
+        assertTrue("profile system server not enabled", res != null && res.equals("true"));
+    }
+
+    private void forceSaveProfile(String pkg) throws Exception {
+        String pid = mTestDevice.executeShellCommand("pidof " + pkg).trim();
+        assertTrue("Invalid pid " + pid, pid.length() > 0);
+        String res = mTestDevice.executeShellCommand("kill -s SIGUSR1 " + pid).trim();
+        assertTrue("kill SIGUSR1: " + res, res.length() == 0);
+    }
+
+    @Test
+    public void testSystemServerProfile() throws Exception {
+        // Trunacte the profile before force it to be saved to prevent previous profiles
+        // causing the test to pass.
+        String res;
+        res = mTestDevice.executeShellCommand("truncate -s 0 " + SYSTEM_SERVER_PROFILE).trim();
+        assertTrue(res, res.length() == 0);
+        // Force save profiles in case the system just started.
+        Thread.sleep(1000);
+        forceSaveProfile("system_server");
+        Thread.sleep(2000);
+        // Validate that the profile is non empty.
+        res = mTestDevice.executeShellCommand("profman --dump-only --profile-file="
+                + SYSTEM_SERVER_PROFILE);
+        boolean sawFramework = false;
+        boolean sawServices = false;
+        for (String line : res.split("\n")) {
+            if (line.contains("framework.jar")) {
+                sawFramework = true;
+            } else if (line.contains("services.jar")) {
+                sawServices = true;
+            }
+        }
+        assertTrue("Did not see framework.jar in " + res, sawFramework);
+        assertTrue("Did not see services.jar in " + res, sawServices);
+    }
+}
diff --git a/tests/PackageWatchdog/Android.bp b/tests/PackageWatchdog/Android.bp
index 88d92c4..0b75039 100644
--- a/tests/PackageWatchdog/Android.bp
+++ b/tests/PackageWatchdog/Android.bp
@@ -23,6 +23,7 @@
         "androidx.test.rules",
         "services.core",
         "services.net",
+        "truth-prebuilt",
     ],
     libs: ["android.test.runner"],
     jni_libs: [
diff --git a/tests/PackageWatchdog/src/com/android/server/PackageWatchdogTest.java b/tests/PackageWatchdog/src/com/android/server/PackageWatchdogTest.java
index 6c05bb8..9a60330 100644
--- a/tests/PackageWatchdog/src/com/android/server/PackageWatchdogTest.java
+++ b/tests/PackageWatchdog/src/com/android/server/PackageWatchdogTest.java
@@ -18,9 +18,8 @@
 
 import static android.service.watchdog.ExplicitHealthCheckService.PackageConfig;
 
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
+import static com.google.common.truth.Truth.assertThat;
+
 import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyString;
@@ -65,7 +64,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.function.Consumer;
 
-// TODO: Use Truth in tests.
 /**
  * Test PackageWatchdog.
  */
@@ -119,13 +117,12 @@
         TestObserver observer = new TestObserver(OBSERVER_NAME_1);
 
         watchdog.startObservingHealth(observer, Arrays.asList(APP_A), SHORT_DURATION);
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // The failed packages should be the same as the registered ones to ensure registration is
         // done successfully
-        assertEquals(1, observer.mHealthCheckFailedPackages.size());
-        assertTrue(observer.mHealthCheckFailedPackages.contains(APP_A));
+        assertThat(observer.mHealthCheckFailedPackages).containsExactly(APP_A);
     }
 
     @Test
@@ -136,17 +133,14 @@
 
         watchdog.startObservingHealth(observer1, Arrays.asList(APP_A), SHORT_DURATION);
         watchdog.startObservingHealth(observer2, Arrays.asList(APP_A, APP_B), SHORT_DURATION);
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE),
-                new VersionedPackage(APP_B, VERSION_CODE)));
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE),
+                        new VersionedPackage(APP_B, VERSION_CODE)));
 
         // The failed packages should be the same as the registered ones to ensure registration is
         // done successfully
-        assertEquals(1, observer1.mHealthCheckFailedPackages.size());
-        assertEquals(2, observer2.mHealthCheckFailedPackages.size());
-        assertTrue(observer1.mHealthCheckFailedPackages.contains(APP_A));
-        assertTrue(observer2.mHealthCheckFailedPackages.contains(APP_A));
-        assertTrue(observer2.mHealthCheckFailedPackages.contains(APP_B));
+        assertThat(observer1.mHealthCheckFailedPackages).containsExactly(APP_A);
+        assertThat(observer2.mHealthCheckFailedPackages).containsExactly(APP_A, APP_B);
     }
 
     @Test
@@ -156,11 +150,11 @@
 
         watchdog.startObservingHealth(observer, Arrays.asList(APP_A), SHORT_DURATION);
         watchdog.unregisterHealthObserver(observer);
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // We should have no failed packages to ensure unregistration is done successfully
-        assertEquals(0, observer.mHealthCheckFailedPackages.size());
+        assertThat(observer.mHealthCheckFailedPackages).isEmpty();
     }
 
     @Test
@@ -172,13 +166,13 @@
         watchdog.startObservingHealth(observer1, Arrays.asList(APP_A), SHORT_DURATION);
         watchdog.startObservingHealth(observer2, Arrays.asList(APP_A), SHORT_DURATION);
         watchdog.unregisterHealthObserver(observer2);
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // observer1 should receive failed packages as intended.
-        assertEquals(1, observer1.mHealthCheckFailedPackages.size());
+        assertThat(observer1.mHealthCheckFailedPackages).containsExactly(APP_A);
         // observer2 should have no failed packages to ensure unregistration is done successfully
-        assertEquals(0, observer2.mHealthCheckFailedPackages.size());
+        assertThat(observer2.mHealthCheckFailedPackages).isEmpty();
     }
 
     @Test
@@ -188,11 +182,11 @@
 
         watchdog.startObservingHealth(observer, Arrays.asList(APP_A), SHORT_DURATION);
         moveTimeForwardAndDispatch(SHORT_DURATION);
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // We should have no failed packages for the fatal failure is raised after expiration
-        assertEquals(0, observer.mHealthCheckFailedPackages.size());
+        assertThat(observer.mHealthCheckFailedPackages).isEmpty();
     }
 
     @Test
@@ -204,13 +198,13 @@
         watchdog.startObservingHealth(observer1, Arrays.asList(APP_A), SHORT_DURATION);
         watchdog.startObservingHealth(observer2, Arrays.asList(APP_A), LONG_DURATION);
         moveTimeForwardAndDispatch(SHORT_DURATION);
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // We should have no failed packages for the fatal failure is raised after expiration
-        assertEquals(0, observer1.mHealthCheckFailedPackages.size());
+        assertThat(observer1.mHealthCheckFailedPackages).isEmpty();
         // We should have failed packages since observer2 hasn't expired
-        assertEquals(1, observer2.mHealthCheckFailedPackages.size());
+        assertThat(observer2.mHealthCheckFailedPackages).containsExactly(APP_A);
     }
 
     /** Observing already observed package extends the observation time. */
@@ -231,12 +225,11 @@
         // Then advance time such that it should have expired were it not for the second observation
         moveTimeForwardAndDispatch((SHORT_DURATION / 2) + 1);
 
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // Verify that we receive failed packages as expected for APP_A not expired
-        assertEquals(1, observer.mHealthCheckFailedPackages.size());
-        assertTrue(observer.mHealthCheckFailedPackages.contains(APP_A));
+        assertThat(observer.mHealthCheckFailedPackages).containsExactly(APP_A);
     }
 
     /**
@@ -257,17 +250,14 @@
         // Then resume observer1 and observer2
         watchdog2.registerHealthObserver(observer1);
         watchdog2.registerHealthObserver(observer2);
-        raiseFatalFailure(watchdog2, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE),
-                new VersionedPackage(APP_B, VERSION_CODE)));
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog2,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE),
+                        new VersionedPackage(APP_B, VERSION_CODE)));
 
         // We should receive failed packages as expected to ensure observers are persisted and
         // resumed correctly
-        assertEquals(1, observer1.mHealthCheckFailedPackages.size());
-        assertEquals(2, observer2.mHealthCheckFailedPackages.size());
-        assertTrue(observer1.mHealthCheckFailedPackages.contains(APP_A));
-        assertTrue(observer1.mHealthCheckFailedPackages.contains(APP_A));
-        assertTrue(observer2.mHealthCheckFailedPackages.contains(APP_B));
+        assertThat(observer1.mHealthCheckFailedPackages).containsExactly(APP_A);
+        assertThat(observer2.mHealthCheckFailedPackages).containsExactly(APP_A, APP_B);
     }
 
     /**
@@ -291,8 +281,8 @@
         mTestLooper.dispatchAll();
 
         // Verify that observers are not notified
-        assertEquals(0, observer1.mMitigatedPackages.size());
-        assertEquals(0, observer2.mMitigatedPackages.size());
+        assertThat(observer1.mHealthCheckFailedPackages).isEmpty();
+        assertThat(observer2.mHealthCheckFailedPackages).isEmpty();
     }
 
     /**
@@ -310,14 +300,12 @@
         watchdog.startObservingHealth(observer1, Arrays.asList(APP_B), SHORT_DURATION);
 
         // Then fail APP_C (not observed) above the threshold
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_C, VERSION_CODE)));
-
-        // Run handler so package failures are dispatched to observers
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_C, VERSION_CODE)));
 
         // Verify that observers are not notified
-        assertEquals(0, observer1.mMitigatedPackages.size());
-        assertEquals(0, observer2.mMitigatedPackages.size());
+        assertThat(observer1.mHealthCheckFailedPackages).isEmpty();
+        assertThat(observer2.mHealthCheckFailedPackages).isEmpty();
     }
 
     /**
@@ -342,14 +330,11 @@
         watchdog.startObservingHealth(observer, Arrays.asList(APP_A), SHORT_DURATION);
 
         // Then fail APP_A (different version) above the threshold
-        raiseFatalFailure(watchdog,
+        raiseFatalFailureAndDispatch(watchdog,
                 Arrays.asList(new VersionedPackage(APP_A, differentVersionCode)));
 
-        // Run handler so package failures are dispatched to observers
-        mTestLooper.dispatchAll();
-
         // Verify that observers are not notified
-        assertEquals(0, observer.mMitigatedPackages.size());
+        assertThat(observer.mHealthCheckFailedPackages).isEmpty();
     }
 
 
@@ -379,13 +364,11 @@
                 SHORT_DURATION);
 
         // Then fail all apps above the threshold
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE),
-                new VersionedPackage(APP_B, VERSION_CODE),
-                new VersionedPackage(APP_C, VERSION_CODE),
-                new VersionedPackage(APP_D, VERSION_CODE)));
-
-        // Run handler so package failures are dispatched to observers
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE),
+                        new VersionedPackage(APP_B, VERSION_CODE),
+                        new VersionedPackage(APP_C, VERSION_CODE),
+                        new VersionedPackage(APP_D, VERSION_CODE)));
 
         // Verify least impact observers are notifed of package failures
         List<String> observerNonePackages = observerNone.mMitigatedPackages;
@@ -394,16 +377,13 @@
         List<String> observerLowPackages = observerLow.mMitigatedPackages;
 
         // APP_D failure observed by only observerNone is not caught cos its impact is none
-        assertEquals(0, observerNonePackages.size());
+        assertThat(observerNonePackages).isEmpty();
         // APP_C failure is caught by observerHigh cos it's the lowest impact observer
-        assertEquals(1, observerHighPackages.size());
-        assertEquals(APP_C, observerHighPackages.get(0));
+        assertThat(observerHighPackages).containsExactly(APP_C);
         // APP_B failure is caught by observerMid cos it's the lowest impact observer
-        assertEquals(1, observerMidPackages.size());
-        assertEquals(APP_B, observerMidPackages.get(0));
+        assertThat(observerMidPackages).containsExactly(APP_B);
         // APP_A failure is caught by observerLow cos it's the lowest impact observer
-        assertEquals(1, observerLowPackages.size());
-        assertEquals(APP_A, observerLowPackages.get(0));
+        assertThat(observerLowPackages).containsExactly(APP_A);
     }
 
     /**
@@ -430,14 +410,12 @@
         watchdog.startObservingHealth(observerSecond, Arrays.asList(APP_A), LONG_DURATION);
 
         // Then fail APP_A above the threshold
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        // Run handler so package failures are dispatched to observers
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // Verify only observerFirst is notifed
-        assertEquals(1, observerFirst.mMitigatedPackages.size());
-        assertEquals(APP_A, observerFirst.mMitigatedPackages.get(0));
-        assertEquals(0, observerSecond.mMitigatedPackages.size());
+        assertThat(observerFirst.mMitigatedPackages).containsExactly(APP_A);
+        assertThat(observerSecond.mMitigatedPackages).isEmpty();
 
         // After observerFirst handles failure, next action it has is high impact
         observerFirst.mImpact = PackageHealthObserverImpact.USER_IMPACT_HIGH;
@@ -445,14 +423,12 @@
         observerSecond.mMitigatedPackages.clear();
 
         // Then fail APP_A again above the threshold
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        // Run handler so package failures are dispatched to observers
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // Verify only observerSecond is notifed cos it has least impact
-        assertEquals(1, observerSecond.mMitigatedPackages.size());
-        assertEquals(APP_A, observerSecond.mMitigatedPackages.get(0));
-        assertEquals(0, observerFirst.mMitigatedPackages.size());
+        assertThat(observerSecond.mMitigatedPackages).containsExactly(APP_A);
+        assertThat(observerFirst.mMitigatedPackages).isEmpty();
 
         // After observerSecond handles failure, it has no further actions
         observerSecond.mImpact = PackageHealthObserverImpact.USER_IMPACT_NONE;
@@ -460,14 +436,12 @@
         observerSecond.mMitigatedPackages.clear();
 
         // Then fail APP_A again above the threshold
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        // Run handler so package failures are dispatched to observers
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // Verify only observerFirst is notifed cos it has the only action
-        assertEquals(1, observerFirst.mMitigatedPackages.size());
-        assertEquals(APP_A, observerFirst.mMitigatedPackages.get(0));
-        assertEquals(0, observerSecond.mMitigatedPackages.size());
+        assertThat(observerFirst.mMitigatedPackages).containsExactly(APP_A);
+        assertThat(observerSecond.mMitigatedPackages).isEmpty();
 
         // After observerFirst handles failure, it too has no further actions
         observerFirst.mImpact = PackageHealthObserverImpact.USER_IMPACT_NONE;
@@ -475,13 +449,12 @@
         observerSecond.mMitigatedPackages.clear();
 
         // Then fail APP_A again above the threshold
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-        // Run handler so package failures are dispatched to observers
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // Verify no observer is notified cos no actions left
-        assertEquals(0, observerFirst.mMitigatedPackages.size());
-        assertEquals(0, observerSecond.mMitigatedPackages.size());
+        assertThat(observerFirst.mMitigatedPackages).isEmpty();
+        assertThat(observerSecond.mMitigatedPackages).isEmpty();
     }
 
     /**
@@ -500,15 +473,12 @@
         watchdog.startObservingHealth(observer1, Arrays.asList(APP_A), SHORT_DURATION);
 
         // Then fail APP_A above the threshold
-        raiseFatalFailure(watchdog, Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
-
-        // Run handler so package failures are dispatched to observers
-        mTestLooper.dispatchAll();
+        raiseFatalFailureAndDispatch(watchdog,
+                Arrays.asList(new VersionedPackage(APP_A, VERSION_CODE)));
 
         // Verify only one observer is notifed
-        assertEquals(1, observer1.mMitigatedPackages.size());
-        assertEquals(APP_A, observer1.mMitigatedPackages.get(0));
-        assertEquals(0, observer2.mMitigatedPackages.size());
+        assertThat(observer1.mMitigatedPackages).containsExactly(APP_A);
+        assertThat(observer2.mMitigatedPackages).isEmpty();
     }
 
     /**
@@ -537,9 +507,7 @@
 
         // Verify we requested health checks for APP_A and APP_B
         List<String> requestedPackages = controller.getRequestedPackages();
-        assertEquals(2, requestedPackages.size());
-        assertEquals(APP_A, requestedPackages.get(0));
-        assertEquals(APP_B, requestedPackages.get(1));
+        assertThat(requestedPackages).containsExactly(APP_A, APP_B);
 
         // Then health check passed for APP_A (observer1 is aware)
         controller.setPackagePassed(APP_A);
@@ -554,18 +522,16 @@
         moveTimeForwardAndDispatch(SHORT_DURATION);
 
         // Verify we cancelled all requests on expiry
-        assertEquals(0, controller.getRequestedPackages().size());
+        assertThat(controller.getRequestedPackages()).isEmpty();
 
         // Verify observer1 is not notified
-        assertEquals(0, observer1.mMitigatedPackages.size());
+        assertThat(observer1.mMitigatedPackages).isEmpty();
 
         // Verify observer2 is notifed because health checks for APP_B never passed
-        assertEquals(1, observer2.mMitigatedPackages.size());
-        assertEquals(APP_B, observer2.mMitigatedPackages.get(0));
+        assertThat(observer2.mMitigatedPackages).containsExactly(APP_B);
 
         // Verify observer3 is notifed because health checks for APP_A did not pass before expiry
-        assertEquals(1, observer3.mMitigatedPackages.size());
-        assertEquals(APP_A, observer3.mMitigatedPackages.get(0));
+        assertThat(observer3.mMitigatedPackages).containsExactly(APP_A);
     }
 
     /**
@@ -592,9 +558,7 @@
 
         // Verify we requested health checks for APP_A and APP_B
         List<String> requestedPackages = controller.getRequestedPackages();
-        assertEquals(2, requestedPackages.size());
-        assertEquals(APP_A, requestedPackages.get(0));
-        assertEquals(APP_B, requestedPackages.get(1));
+        assertThat(requestedPackages).containsExactly(APP_A, APP_B);
 
         // Disable explicit health checks (marks APP_A and APP_B as passed)
         setExplicitHealthCheckEnabled(false);
@@ -603,13 +567,13 @@
         mTestLooper.dispatchAll();
 
         // Verify all checks are cancelled
-        assertEquals(0, controller.getRequestedPackages().size());
+        assertThat(controller.getRequestedPackages()).isEmpty();
 
         // Then expire APP_A
         moveTimeForwardAndDispatch(SHORT_DURATION);
 
         // Verify APP_A is not failed (APP_B) is not expired yet
-        assertEquals(0, observer.mMitigatedPackages.size());
+        assertThat(observer.mMitigatedPackages).isEmpty();
 
         // Re-enable explicit health checks
         setExplicitHealthCheckEnabled(true);
@@ -618,7 +582,7 @@
         mTestLooper.dispatchAll();
 
         // Verify no requests are made cos APP_A is expired and APP_B was marked as passed
-        assertEquals(0, controller.getRequestedPackages().size());
+        assertThat(controller.getRequestedPackages()).isEmpty();
 
         // Then set new supported packages
         controller.setSupportedPackages(Arrays.asList(APP_C));
@@ -630,15 +594,13 @@
 
         // Verify requests are only made for APP_C
         requestedPackages = controller.getRequestedPackages();
-        assertEquals(1, requestedPackages.size());
-        assertEquals(APP_C, requestedPackages.get(0));
+        assertThat(requestedPackages).containsExactly(APP_C);
 
         // Then expire APP_A and APP_C
         moveTimeForwardAndDispatch(SHORT_DURATION);
 
         // Verify only APP_C is failed because explicit health checks was not supported for APP_A
-        assertEquals(1, observer.mMitigatedPackages.size());
-        assertEquals(APP_C, observer.mMitigatedPackages.get(0));
+        assertThat(observer.mMitigatedPackages).containsExactly(APP_C);
     }
 
     /**
@@ -662,8 +624,7 @@
         moveTimeForwardAndDispatch(SHORT_DURATION);
 
         // Verify that health check is failed
-        assertEquals(1, observer.mMitigatedPackages.size());
-        assertEquals(APP_A, observer.mMitigatedPackages.get(0));
+        assertThat(observer.mMitigatedPackages).containsExactly(APP_A);
 
         // Then clear failed packages and start observing a random package so requests are synced
         // and PackageWatchdog#onSupportedPackages is called and APP_A has a chance to fail again
@@ -672,7 +633,7 @@
         watchdog.startObservingHealth(observer, Arrays.asList(APP_B), LONG_DURATION);
 
         // Verify that health check failure is not notified again
-        assertTrue(observer.mMitigatedPackages.isEmpty());
+        assertThat(observer.mMitigatedPackages).isEmpty();
     }
 
     /** Tests {@link MonitoredPackage} health check state transitions. */
@@ -688,36 +649,38 @@
 
         // Verify transition: inactive -> active -> passed
         // Verify initially inactive
-        assertEquals(HealthCheckState.INACTIVE, m1.getHealthCheckStateLocked());
+        assertThat(m1.getHealthCheckStateLocked()).isEqualTo(HealthCheckState.INACTIVE);
         // Verify still inactive, until we #setHealthCheckActiveLocked
-        assertEquals(HealthCheckState.INACTIVE, m1.handleElapsedTimeLocked(SHORT_DURATION));
+        assertThat(m1.handleElapsedTimeLocked(SHORT_DURATION)).isEqualTo(HealthCheckState.INACTIVE);
         // Verify now active
-        assertEquals(HealthCheckState.ACTIVE, m1.setHealthCheckActiveLocked(SHORT_DURATION));
+        assertThat(m1.setHealthCheckActiveLocked(SHORT_DURATION)).isEqualTo(
+                HealthCheckState.ACTIVE);
         // Verify now passed
-        assertEquals(HealthCheckState.PASSED, m1.tryPassHealthCheckLocked());
+        assertThat(m1.tryPassHealthCheckLocked()).isEqualTo(HealthCheckState.PASSED);
 
         // Verify transition: inactive -> active -> failed
         // Verify initially inactive
-        assertEquals(HealthCheckState.INACTIVE, m2.getHealthCheckStateLocked());
+        assertThat(m2.getHealthCheckStateLocked()).isEqualTo(HealthCheckState.INACTIVE);
         // Verify now active
-        assertEquals(HealthCheckState.ACTIVE, m2.setHealthCheckActiveLocked(SHORT_DURATION));
+        assertThat(m2.setHealthCheckActiveLocked(SHORT_DURATION)).isEqualTo(
+                HealthCheckState.ACTIVE);
         // Verify now failed
-        assertEquals(HealthCheckState.FAILED, m2.handleElapsedTimeLocked(SHORT_DURATION));
+        assertThat(m2.handleElapsedTimeLocked(SHORT_DURATION)).isEqualTo(HealthCheckState.FAILED);
 
         // Verify transition: inactive -> failed
         // Verify initially inactive
-        assertEquals(HealthCheckState.INACTIVE, m3.getHealthCheckStateLocked());
+        assertThat(m3.getHealthCheckStateLocked()).isEqualTo(HealthCheckState.INACTIVE);
         // Verify now failed because package expired
-        assertEquals(HealthCheckState.FAILED, m3.handleElapsedTimeLocked(LONG_DURATION));
+        assertThat(m3.handleElapsedTimeLocked(LONG_DURATION)).isEqualTo(HealthCheckState.FAILED);
         // Verify remains failed even when asked to pass
-        assertEquals(HealthCheckState.FAILED, m3.tryPassHealthCheckLocked());
+        assertThat(m3.tryPassHealthCheckLocked()).isEqualTo(HealthCheckState.FAILED);
 
         // Verify transition: passed
-        assertEquals(HealthCheckState.PASSED, m4.getHealthCheckStateLocked());
+        assertThat(m4.getHealthCheckStateLocked()).isEqualTo(HealthCheckState.PASSED);
         // Verify remains passed even if health check fails
-        assertEquals(HealthCheckState.PASSED, m4.handleElapsedTimeLocked(SHORT_DURATION));
+        assertThat(m4.handleElapsedTimeLocked(SHORT_DURATION)).isEqualTo(HealthCheckState.PASSED);
         // Verify remains passed even if package expires
-        assertEquals(HealthCheckState.PASSED, m4.handleElapsedTimeLocked(LONG_DURATION));
+        assertThat(m4.handleElapsedTimeLocked(LONG_DURATION)).isEqualTo(HealthCheckState.PASSED);
     }
 
     @Test
@@ -736,8 +699,7 @@
         mTestLooper.dispatchAll();
 
         // Verify the NetworkStack observer is notified
-        assertEquals(1, observer.mMitigatedPackages.size());
-        assertEquals(APP_A, observer.mMitigatedPackages.get(0));
+        assertThat(observer.mMitigatedPackages).containsExactly(APP_A);
     }
 
     private void adoptShellPermissions(String... permissions) {
@@ -773,10 +735,12 @@
     }
 
     /** Trigger package failures above the threshold. */
-    private void raiseFatalFailure(PackageWatchdog watchdog, List<VersionedPackage> packages) {
+    private void raiseFatalFailureAndDispatch(PackageWatchdog watchdog,
+            List<VersionedPackage> packages) {
         for (int i = 0; i < watchdog.getTriggerFailureCount(); i++) {
             watchdog.onPackageFailure(packages);
         }
+        mTestLooper.dispatchAll();
     }
 
     private PackageWatchdog createWatchdog() {
@@ -791,13 +755,13 @@
                 new PackageWatchdog(mSpyContext, policyFile, handler, handler, controller,
                         mConnectivityModuleConnector, mTestClock);
         // Verify controller is not automatically started
-        assertFalse(controller.mIsEnabled);
+        assertThat(controller.mIsEnabled).isFalse();
         if (withPackagesReady) {
             // Only capture the NetworkStack callback for the latest registered watchdog
             reset(mConnectivityModuleConnector);
             watchdog.onPackagesReady();
             // Verify controller by default is started when packages are ready
-            assertTrue(controller.mIsEnabled);
+            assertThat(controller.mIsEnabled).isTrue();
 
             verify(mConnectivityModuleConnector).registerHealthListener(
                     mConnectivityModuleCallbackCaptor.capture());