[KV] Refactor sendDataToTransport() and finishTask()

Refactor sendDataToTransport():
* Instead of truncating backup data when reading backup data / writing widget
  data, consider the IOException a transport-level failure. Added tests for
  this.
* Extracted a few methods: validateBackupData(), updateFiles(),
  handleTransportStatus().
* Put the if (size > 0) check outside try-catch block.
* Used try-with-resources.
* Clean files in case of transport error, quota exceeded.

Refactor finishTask():
* Extracted triggerTransportInitializationLocked(), assigned queue lock
  to private final var in ctor.
* In triggerTransportInitializationLocked() set the status to T_ERROR if
  we failed to query the name of the transport.
* Tests for TRANSPORT_NOT_INITIALIZED.

General:
* Small refactors in KVBT.
* Small refactors in test.

Bug: 110082831
Bug: 113311470
Test: adb shell bmgr backupnow <kv_package>
Test: 1. adb shell bmgr backupnow <kv_package>
      2. Transport returns T_NOT_INITIALIZED
      3. Make sure PM metadata state file is deleted.
Test: atest KeyValueBackupTaskTest
Change-Id: I8d85c24cba6da4fbaf14234e2ce6d8e0699a3eed
diff --git a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupReporter.java b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupReporter.java
index 8fbca4b..7f5ddc2 100644
--- a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupReporter.java
+++ b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupReporter.java
@@ -26,6 +26,7 @@
 import android.util.EventLog;
 import android.util.Slog;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.server.EventLogTags;
 import com.android.server.backup.BackupManagerService;
 import com.android.server.backup.DataChangedJournal;
@@ -49,10 +50,19 @@
  */
 // TODO: In KeyValueBackupTaskTest, remove direct assertions on logcat, observer or monitor and
 //       verify calls to this object. Add these and more assertions to the test of this class.
-class KeyValueBackupReporter {
-    private static final String TAG = "KeyValueBackupTask";
+@VisibleForTesting
+public class KeyValueBackupReporter {
+    @VisibleForTesting
+    static final String TAG = "KeyValueBackupTask";
     private static final boolean DEBUG = BackupManagerService.DEBUG;
-    private static final boolean MORE_DEBUG = BackupManagerService.MORE_DEBUG || true;
+    @VisibleForTesting
+    static final boolean MORE_DEBUG = BackupManagerService.MORE_DEBUG || true;
+
+    static void onNewThread(String threadName) {
+        if (DEBUG) {
+            Slog.d(TAG, "Spinning thread " + threadName);
+        }
+    }
 
     private final BackupManagerService mBackupManagerService;
     private final IBackupObserver mObserver;
@@ -61,7 +71,7 @@
     KeyValueBackupReporter(
             BackupManagerService backupManagerService,
             IBackupObserver observer,
-            IBackupManagerMonitor monitor) {
+            @Nullable IBackupManagerMonitor monitor) {
         mBackupManagerService = backupManagerService;
         mObserver = observer;
         mMonitor = monitor;
@@ -73,6 +83,10 @@
         return mMonitor;
     }
 
+    IBackupObserver getObserver() {
+        return mObserver;
+    }
+
     void onSkipBackup() {
         if (DEBUG) {
             Slog.d(TAG, "Skipping backup since one is already in progress");
@@ -237,10 +251,6 @@
         }
     }
 
-    void onTruncateDataError() {
-        Slog.w(TAG, "Unable to roll back");
-    }
-
     void onSendDataToTransport(String packageName) {
         if (MORE_DEBUG) {
             Slog.v(TAG, "Sending non-empty data to transport for " + packageName);
@@ -357,7 +367,7 @@
         return (packageInfo != null) ? packageInfo.packageName : "no_package_yet";
     }
 
-    void onRevertBackup() {
+    void onRevertTask() {
         if (MORE_DEBUG) {
             Slog.i(TAG, "Reverting backup queue, re-staging everything");
         }
@@ -391,6 +401,10 @@
         Slog.w(TAG, "Failed to query transport name for pending init: " + e);
     }
 
+    /**
+     * This is a bit different from {@link #onTaskFinished()}, it's only called if there is no
+     * full-backup requests associated with the key-value task.
+     */
     void onBackupFinished(int status) {
         BackupObserverUtils.sendBackupFinished(mObserver, status);
     }
@@ -399,7 +413,7 @@
         Slog.d(TAG, "Starting full backups for: " + pendingFullBackups);
     }
 
-    void onKeyValueBackupFinished() {
+    void onTaskFinished() {
         Slog.i(TAG, "K/V backup pass finished");
     }
 }
diff --git a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
index 91af6f1..7d035cb 100644
--- a/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
+++ b/services/backup/java/com/android/server/backup/keyvalue/KeyValueBackupTask.java
@@ -24,7 +24,6 @@
 import static com.android.server.backup.BackupManagerService.KEY_WIDGET_STATE;
 import static com.android.server.backup.BackupManagerService.OP_PENDING;
 import static com.android.server.backup.BackupManagerService.OP_TYPE_BACKUP;
-import static com.android.server.backup.BackupManagerService.PACKAGE_MANAGER_SENTINEL;
 
 import android.annotation.Nullable;
 import android.app.ApplicationThreadConstants;
@@ -48,11 +47,9 @@
 import android.os.SELinux;
 import android.os.UserHandle;
 import android.os.WorkSource;
-import android.system.ErrnoException;
-import android.system.Os;
 import android.util.Pair;
-import android.util.Slog;
 
+import com.android.internal.annotations.GuardedBy;
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.backup.IBackupTransport;
 import com.android.internal.util.Preconditions;
@@ -172,9 +169,6 @@
 // TODO: Consider having the caller responsible for some clean-up (like resetting state)
 // TODO: Distinguish between cancel and time-out where possible for logging/monitoring/observing
 public class KeyValueBackupTask implements BackupRestoreTask, Runnable {
-    private static final String TAG = "KeyValueBackupTask";
-    private static final boolean DEBUG = BackupManagerService.DEBUG;
-    private static final boolean MORE_DEBUG = BackupManagerService.MORE_DEBUG || false;
     private static final int THREAD_PRIORITY = Process.THREAD_PRIORITY_BACKGROUND;
     private static final AtomicInteger THREAD_COUNT = new AtomicInteger();
     private static final String BLANK_STATE_FILE_NAME = "blank_state";
@@ -218,6 +212,8 @@
             List<String> pendingFullBackups,
             boolean userInitiated,
             boolean nonIncremental) {
+        KeyValueBackupReporter reporter =
+                new KeyValueBackupReporter(backupManagerService, observer, monitor);
         KeyValueBackupTask task =
                 new KeyValueBackupTask(
                         backupManagerService,
@@ -225,17 +221,14 @@
                         transportDirName,
                         queue,
                         dataChangedJournal,
-                        observer,
-                        monitor,
+                        reporter,
                         listener,
                         pendingFullBackups,
                         userInitiated,
                         nonIncremental);
         Thread thread = new Thread(task, "key-value-backup-" + THREAD_COUNT.incrementAndGet());
-        if (DEBUG) {
-            Slog.d(TAG, "Spinning thread " + thread.getName());
-        }
         thread.start();
+        KeyValueBackupReporter.onNewThread(thread.getName());
         return task;
     }
 
@@ -244,16 +237,17 @@
     private final TransportManager mTransportManager;
     private final TransportClient mTransportClient;
     private final BackupAgentTimeoutParameters mAgentTimeoutParameters;
-    private final IBackupObserver mObserver;
     private final KeyValueBackupReporter mReporter;
     private final OnTaskFinishedListener mTaskFinishedListener;
     private final boolean mUserInitiated;
     private final boolean mNonIncremental;
     private final int mCurrentOpToken;
-    private final File mStateDir;
+    private final File mStateDirectory;
+    private final File mDataDirectory;
     private final List<String> mOriginalQueue;
     private final List<String> mQueue;
     private final List<String> mPendingFullBackups;
+    private final Object mQueueLock;
     @Nullable private final DataChangedJournal mJournal;
     @Nullable private PerformFullTransportBackupTask mFullBackupTask;
 
@@ -296,8 +290,7 @@
             String transportDirName,
             List<String> queue,
             @Nullable DataChangedJournal journal,
-            IBackupObserver observer,
-            @Nullable IBackupManagerMonitor monitor,
+            KeyValueBackupReporter reporter,
             OnTaskFinishedListener taskFinishedListener,
             List<String> pendingFullBackups,
             boolean userInitiated,
@@ -310,8 +303,7 @@
         // We need to retain the original queue contents in case of transport failure
         mQueue = new ArrayList<>(queue);
         mJournal = journal;
-        mObserver = observer;
-        mReporter = new KeyValueBackupReporter(backupManagerService, observer, monitor);
+        mReporter = reporter;
         mTaskFinishedListener = taskFinishedListener;
         mPendingFullBackups = pendingFullBackups;
         mUserInitiated = userInitiated;
@@ -320,8 +312,10 @@
                 Preconditions.checkNotNull(
                         backupManagerService.getAgentTimeoutParameters(),
                         "Timeout parameters cannot be null");
-        mStateDir = new File(backupManagerService.getBaseStateDir(), transportDirName);
+        mStateDirectory = new File(backupManagerService.getBaseStateDir(), transportDirName);
+        mDataDirectory = mBackupManagerService.getDataDir();
         mCurrentOpToken = backupManagerService.generateRandomIntegerToken();
+        mQueueLock = mBackupManagerService.getQueueLock();
     }
 
     private void registerTask() {
@@ -337,7 +331,7 @@
     public void run() {
         Process.setThreadPriority(THREAD_PRIORITY);
 
-        BackupState state = startBackup();
+        BackupState state = startTask();
         while (state == BackupState.RUNNING_QUEUE || state == BackupState.BACKUP_PM) {
             if (mCancelled) {
                 state = BackupState.CANCELLED;
@@ -355,7 +349,7 @@
                     break;
             }
         }
-        finishBackup();
+        finishTask();
     }
 
     private BackupState handleAgentResult(RemoteResult result) {
@@ -383,7 +377,7 @@
     @Override
     public void operationComplete(long unusedResult) {}
 
-    private BackupState startBackup() {
+    private BackupState startTask() {
         synchronized (mBackupManagerService.getCurrentOpLock()) {
             if (mBackupManagerService.isBackupOperationInProgress()) {
                 mReporter.onSkipBackup();
@@ -401,7 +395,7 @@
                         /* updateSchedule */ false,
                         /* runningJob */ null,
                         new CountDownLatch(1),
-                        mObserver,
+                        mReporter.getObserver(),
                         mReporter.getMonitor(),
                         mTaskFinishedListener,
                         mUserInitiated);
@@ -420,16 +414,16 @@
         boolean backupPm = mQueue.remove(PM_PACKAGE) || !mNonIncremental;
 
         mReporter.onQueueReady(mQueue);
-        File pmState = new File(mStateDir, PM_PACKAGE);
+        File pmState = new File(mStateDirectory, PM_PACKAGE);
         try {
-            IBackupTransport transport = mTransportClient.connectOrThrow("KVBT.startBackup()");
+            IBackupTransport transport = mTransportClient.connectOrThrow("KVBT.startTask()");
             String transportName = transport.name();
             mReporter.onTransportReady(transportName);
 
             // If we haven't stored PM metadata yet, we must initialize the transport.
             if (pmState.length() <= 0) {
                 mReporter.onInitializeTransport(transportName);
-                mBackupManagerService.resetBackupState(mStateDir);
+                mBackupManagerService.resetBackupState(mStateDirectory);
                 mStatus = transport.initializeDevice();
                 mReporter.onTransportInitialized(mStatus);
             }
@@ -439,7 +433,7 @@
         }
 
         if (mStatus != BackupTransport.TRANSPORT_OK) {
-            mBackupManagerService.resetBackupState(mStateDir);
+            mBackupManagerService.resetBackupState(mStateDirectory);
             return BackupState.FINAL;
         }
 
@@ -471,7 +465,7 @@
         }
 
         if (mStatus != BackupTransport.TRANSPORT_OK) {
-            mBackupManagerService.resetBackupState(mStateDir);
+            mBackupManagerService.resetBackupState(mStateDirectory);
             return BackupState.FINAL;
         }
 
@@ -566,7 +560,7 @@
             }
 
             // Transport-level failure, re-enqueue everything.
-            revertBackup();
+            revertTask();
             return Pair.create(BackupState.FINAL, null);
         }
 
@@ -574,7 +568,7 @@
         return Pair.create(null, agentResult);
     }
 
-    private void finishBackup() {
+    private void finishTask() {
         // Mark packages that we couldn't backup as pending backup.
         for (String packageName : mQueue) {
             mBackupManagerService.dataChangedImpl(packageName);
@@ -586,7 +580,7 @@
             mReporter.onJournalDeleteFailed(mJournal);
         }
 
-        String callerLogString = "KVBT.finishBackup()";
+        String callerLogString = "KVBT.finishTask()";
 
         // If we succeeded and this is the first time we've done a backup, we can record the current
         // backup dataset token.
@@ -602,23 +596,16 @@
             }
         }
 
-        synchronized (mBackupManagerService.getQueueLock()) {
+        synchronized (mQueueLock) {
             mBackupManagerService.setBackupRunning(false);
             if (mStatus == BackupTransport.TRANSPORT_NOT_INITIALIZED) {
                 mReporter.onTransportNotInitialized();
-                try {
-                    IBackupTransport transport = mTransportClient.connectOrThrow(callerLogString);
-                    mBackupManagerService.getPendingInits().add(transport.name());
-                    clearPmMetadata();
-                    mBackupManagerService.backupNow();
-                } catch (Exception e) {
-                    mReporter.onPendingInitializeTransportError(e);
-                }
+                triggerTransportInitializationLocked();
             }
         }
 
         unregisterTask();
-        mReporter.onKeyValueBackupFinished();
+        mReporter.onTaskFinished();
 
         if (mCancelled) {
             // We acknowledge the cancel as soon as we unregister the task, allowing other backups
@@ -663,14 +650,25 @@
         }
     }
 
-    /** Removes PM state, triggering initialization in the next key-value task. */
-    private void clearPmMetadata() {
-        File pmState = new File(mStateDir, PM_PACKAGE);
-        if (pmState.exists()) {
-            pmState.delete();
+    @GuardedBy("mQueueLock")
+    private void triggerTransportInitializationLocked() {
+        try {
+            IBackupTransport transport =
+                    mTransportClient.connectOrThrow("KVBT.triggerTransportInitializationLocked");
+            mBackupManagerService.getPendingInits().add(transport.name());
+            deletePmStateFile();
+            mBackupManagerService.backupNow();
+        } catch (Exception e) {
+            mReporter.onPendingInitializeTransportError(e);
+            mStatus = BackupTransport.TRANSPORT_ERROR;
         }
     }
 
+    /** Removes PM state, triggering initialization in the next key-value task. */
+    private void deletePmStateFile() {
+        new File(mStateDirectory, PM_PACKAGE).delete();
+    }
+
     /**
      * Returns a {@link Pair}. The first of the pair contains the status. In case the status is
      * {@link BackupTransport#TRANSPORT_OK}, the second of the pair contains the agent result,
@@ -679,12 +677,10 @@
     private Pair<Integer, RemoteResult> extractAgentData(String packageName, IBackupAgent agent) {
         mReporter.onExtractAgentData(packageName);
 
-        File blankStateFile = new File(mStateDir, BLANK_STATE_FILE_NAME);
-        mSavedStateFile = new File(mStateDir, packageName);
-        File savedStateFileForAgent = (mNonIncremental) ? blankStateFile : mSavedStateFile;
-        mBackupDataFile =
-                new File(mBackupManagerService.getDataDir(), packageName + STAGING_FILE_SUFFIX);
-        mNewStateFile = new File(mStateDir, packageName + NEW_STATE_FILE_SUFFIX);
+        File blankStateFile = new File(mStateDirectory, BLANK_STATE_FILE_NAME);
+        mSavedStateFile = new File(mStateDirectory, packageName);
+        mBackupDataFile = new File(mDataDirectory, packageName + STAGING_FILE_SUFFIX);
+        mNewStateFile = new File(mStateDirectory, packageName + NEW_STATE_FILE_SUFFIX);
         mReporter.onAgentFilesReady(mBackupDataFile);
 
         mSavedState = null;
@@ -694,6 +690,7 @@
         boolean callingAgent = false;
         final RemoteResult agentResult;
         try {
+            File savedStateFileForAgent = (mNonIncremental) ? blankStateFile : mSavedStateFile;
             // MODE_CREATE to make an empty file if necessary
             mSavedState = ParcelFileDescriptor.open(
                     savedStateFileForAgent, MODE_READ_ONLY | MODE_CREATE);
@@ -737,7 +734,7 @@
         return Pair.create(BackupTransport.TRANSPORT_OK, agentResult);
     }
 
-    private void failAgent(IBackupAgent agent, String message) {
+    private void agentFail(IBackupAgent agent, String message) {
         try {
             agent.fail(message);
         } catch (Exception e) {
@@ -767,7 +764,7 @@
             throws IOException {
         // TODO: http://b/22388012
         byte[] widgetState = AppWidgetBackupBridge.getWidgetState(pkgName, UserHandle.USER_SYSTEM);
-        File widgetFile = new File(mStateDir, pkgName + "_widget");
+        File widgetFile = new File(mStateDirectory, pkgName + "_widget");
         boolean priorStateExists = widgetFile.exists();
         if (!priorStateExists && widgetState == null) {
             return;
@@ -819,154 +816,161 @@
 
         String packageName = mCurrentPackage.packageName;
         ApplicationInfo applicationInfo = mCurrentPackage.applicationInfo;
-        long filePos = mBackupDataFile.length();
-        FileDescriptor fd = mBackupData.getFileDescriptor();
+
         boolean writingWidgetData = false;
         try {
-            // If it's a 3rd party app, crash them if they wrote any protected keys.
-            if (applicationInfo != null &&
-                    (applicationInfo.flags & ApplicationInfo.FLAG_SYSTEM) == 0) {
-                ParcelFileDescriptor readFd =
-                        ParcelFileDescriptor.open(mBackupDataFile, MODE_READ_ONLY);
-                BackupDataInput in = new BackupDataInput(readFd.getFileDescriptor());
-                try {
-                    while (in.readNextHeader()) {
-                        String key = in.getKey();
-                        if (key != null && key.charAt(0) >= 0xff00) {
-                            mReporter.onAgentIllegalKey(mCurrentPackage, key);
-                            failAgent(mAgentBinder, "Illegal backup key: " + key);
-                            errorCleanup();
-                            return BackupState.RUNNING_QUEUE;
-                        }
-                        in.skipEntityData();
-                    }
-                } finally {
-                    readFd.close();
-                }
+            if (!validateBackupData(applicationInfo, mBackupDataFile)) {
+                errorCleanup();
+                return BackupState.RUNNING_QUEUE;
             }
-
             writingWidgetData = true;
-            writeWidgetPayloadIfAppropriate(fd, packageName);
+            writeWidgetPayloadIfAppropriate(mBackupData.getFileDescriptor(), packageName);
         } catch (IOException e) {
             if (writingWidgetData) {
                 mReporter.onWriteWidgetDataError(packageName, e);
             } else {
                 mReporter.onReadAgentDataError(packageName, e);
             }
-            try {
-                Os.ftruncate(fd, filePos);
-            } catch (ErrnoException ee) {
-                mReporter.onTruncateDataError();
-            }
+            revertTask();
+            return BackupState.FINAL;
         }
 
         clearAgentState();
+        boolean nonIncremental = mSavedStateFile.length() == 0;
+        long size = mBackupDataFile.length();
+        if (size > 0) {
+            try (ParcelFileDescriptor backupData =
+                         ParcelFileDescriptor.open(mBackupDataFile, MODE_READ_ONLY)) {
+                IBackupTransport transport =
+                        mTransportClient.connectOrThrow("KVBT.sendDataToTransport()");
+                mReporter.onSendDataToTransport(packageName);
+                int flags = getPerformBackupFlags(mUserInitiated, nonIncremental);
 
-        ParcelFileDescriptor backupData = null;
-        mStatus = BackupTransport.TRANSPORT_OK;
-        long size = 0;
-        try {
-            IBackupTransport transport =
-                    mTransportClient.connectOrThrow("KVBT.sendDataToTransport()");
-            size = mBackupDataFile.length();
-            if (size > 0) {
-                boolean isNonIncremental = mSavedStateFile.length() == 0;
-
-                if (mStatus == BackupTransport.TRANSPORT_OK) {
-                    mReporter.onSendDataToTransport(packageName);
-                    backupData = ParcelFileDescriptor.open(mBackupDataFile, MODE_READ_ONLY);
-                    int userInitiatedFlag =
-                            mUserInitiated ? BackupTransport.FLAG_USER_INITIATED : 0;
-                    int incrementalFlag =
-                            isNonIncremental
-                                    ? BackupTransport.FLAG_NON_INCREMENTAL
-                                    : BackupTransport.FLAG_INCREMENTAL;
-                    int flags = userInitiatedFlag | incrementalFlag;
-
-                    mStatus = transport.performBackup(mCurrentPackage, backupData, flags);
-                }
-
-                if (isNonIncremental
-                        && mStatus == BackupTransport.TRANSPORT_NON_INCREMENTAL_BACKUP_REQUIRED) {
-                    mReporter.onNonIncrementalAndNonIncrementalRequired();
-                    mStatus = BackupTransport.TRANSPORT_ERROR;
-                }
-
+                mStatus = transport.performBackup(mCurrentPackage, backupData, flags);
                 if (mStatus == BackupTransport.TRANSPORT_OK) {
                     mStatus = transport.finishBackup();
                 }
-            } else {
-                mReporter.onEmptyData(mCurrentPackage);
+            } catch (Exception e) {
+                mReporter.onPackageBackupError(packageName, e);
+                mStatus = BackupTransport.TRANSPORT_ERROR;
             }
+        } else {
+            mReporter.onEmptyData(mCurrentPackage);
+            mStatus = BackupTransport.TRANSPORT_OK;
+        }
 
-            if (mStatus == BackupTransport.TRANSPORT_OK) {
+        if (nonIncremental
+                && mStatus == BackupTransport.TRANSPORT_NON_INCREMENTAL_BACKUP_REQUIRED) {
+            mReporter.onNonIncrementalAndNonIncrementalRequired();
+            mStatus = BackupTransport.TRANSPORT_ERROR;
+        }
+
+        updateFiles(mStatus);
+        return handleTransportStatus(mStatus, packageName, size);
+    }
+
+    private void updateFiles(int status) {
+        switch (status) {
+            case BackupTransport.TRANSPORT_OK:
                 mBackupDataFile.delete();
                 mNewStateFile.renameTo(mSavedStateFile);
-                mReporter.onPackageBackupComplete(packageName, size);
-            } else if (mStatus == BackupTransport.TRANSPORT_PACKAGE_REJECTED) {
-                mBackupDataFile.delete();
-                mNewStateFile.delete();
-                mReporter.onPackageBackupRejected(packageName);
-            } else if (mStatus == BackupTransport.TRANSPORT_QUOTA_EXCEEDED) {
-                // TODO: Should reset files like above?
-                mReporter.onPackageBackupQuotaExceeded(packageName);
-            } else if (mStatus == BackupTransport.TRANSPORT_NON_INCREMENTAL_BACKUP_REQUIRED) {
-                mReporter.onPackageBackupNonIncrementalRequired(mCurrentPackage);
-                mBackupDataFile.delete();
+                break;
+            case BackupTransport.TRANSPORT_NON_INCREMENTAL_BACKUP_REQUIRED:
                 mSavedStateFile.delete();
+                mBackupDataFile.delete();
                 mNewStateFile.delete();
+                break;
+            default:
+                // Includes:
+                // * BackupTransport.TRANSPORT_PACKAGE_REJECTED
+                // * BackupTransport.TRANSPORT_QUOTA_EXCEEDED
+                // * BackupTransport.TRANSPORT_ERROR
+                mBackupDataFile.delete();
+                mNewStateFile.delete();
+                break;
 
-                // Immediately retry the package by adding it back to the front of the queue.
-                // We cannot add @pm@ to the queue because we back it up separately at the start.
-                // Below we request PM backup if that is the case.
-                if (!PM_PACKAGE.equals(packageName)) {
-                    mQueue.add(0, packageName);
-                }
-            } else {
-                mReporter.onPackageBackupTransportFailure(packageName);
-            }
-        } catch (Exception e) {
-            mReporter.onPackageBackupError(packageName, e);
-            mStatus = BackupTransport.TRANSPORT_ERROR;
-        } finally {
-            tryCloseFileDescriptor(backupData, "backup data");
         }
+    }
 
-        final BackupState nextState;
-        if (mStatus == BackupTransport.TRANSPORT_OK
-                || mStatus == BackupTransport.TRANSPORT_PACKAGE_REJECTED) {
-            nextState = BackupState.RUNNING_QUEUE;
-
-        } else if (mStatus == BackupTransport.TRANSPORT_NON_INCREMENTAL_BACKUP_REQUIRED) {
-            // We want to immediately retry the current package.
+    private BackupState handleTransportStatus(int status, String packageName, long size) {
+        if (status == BackupTransport.TRANSPORT_OK) {
+            mReporter.onPackageBackupComplete(packageName, size);
+            return BackupState.RUNNING_QUEUE;
+        }
+        if (status == BackupTransport.TRANSPORT_PACKAGE_REJECTED) {
+            mReporter.onPackageBackupRejected(packageName);
+            return BackupState.RUNNING_QUEUE;
+        }
+        if (status == BackupTransport.TRANSPORT_NON_INCREMENTAL_BACKUP_REQUIRED) {
+            mReporter.onPackageBackupNonIncrementalRequired(mCurrentPackage);
+            // Immediately retry the current package.
             if (PM_PACKAGE.equals(packageName)) {
-                nextState = BackupState.BACKUP_PM;
-            } else {
-                // This is an ordinary package so we will have added it back into the queue
-                // above. Thus, we proceed processing the queue.
-                nextState = BackupState.RUNNING_QUEUE;
+                return BackupState.BACKUP_PM;
             }
-
-        } else if (mStatus == BackupTransport.TRANSPORT_QUOTA_EXCEEDED) {
-            if (mAgentBinder != null) {
-                try {
-                    IBackupTransport transport =
-                            mTransportClient.connectOrThrow("KVBT.sendDataToTransport()");
-                    long quota = transport.getBackupQuota(mCurrentPackage.packageName, false);
-                    mAgentBinder.doQuotaExceeded(size, quota);
-                } catch (Exception e) {
-                    mReporter.onAgentDoQuotaExceededError(e);
-                }
-            }
-            nextState = BackupState.RUNNING_QUEUE;
-        } else {
-            // Any other error here indicates a transport-level failure.  That means
-            // we need to halt everything and reschedule everything for next time.
-            revertBackup();
-            nextState = BackupState.FINAL;
+            mQueue.add(0, packageName);
+            return BackupState.RUNNING_QUEUE;
         }
+        if (status == BackupTransport.TRANSPORT_QUOTA_EXCEEDED) {
+            mReporter.onPackageBackupQuotaExceeded(packageName);
+            agentDoQuotaExceeded(mAgentBinder, packageName, size);
+            return BackupState.RUNNING_QUEUE;
+        }
+        // Any other error here indicates a transport-level failure.
+        mReporter.onPackageBackupTransportFailure(packageName);
+        revertTask();
+        return BackupState.FINAL;
+    }
 
-        return nextState;
+    private void agentDoQuotaExceeded(
+            @Nullable IBackupAgent agent, String packageName, long backupDataSize) {
+        if (agent != null) {
+            try {
+                IBackupTransport transport =
+                        mTransportClient.connectOrThrow("KVBT.agentDoQuotaExceeded()");
+                long quota = transport.getBackupQuota(packageName, false);
+                agent.doQuotaExceeded(backupDataSize, quota);
+            } catch (Exception e) {
+                mReporter.onAgentDoQuotaExceededError(e);
+            }
+        }
+    }
+
+    /**
+     * For system apps and pseudo-apps always return {@code true}. For regular apps returns whether
+     * {@code backupDataFile} doesn't have any protected keys.
+     *
+     * <p>If the app has attempted to write any protected keys we also crash them.
+     */
+    private boolean validateBackupData(
+            @Nullable ApplicationInfo applicationInfo, File backupDataFile) throws IOException {
+        if (applicationInfo == null || (applicationInfo.flags & ApplicationInfo.FLAG_SYSTEM) != 0) {
+            // System apps and pseudo-apps can write what they want.
+            return true;
+        }
+        try (ParcelFileDescriptor backupData =
+                     ParcelFileDescriptor.open(backupDataFile, MODE_READ_ONLY)) {
+            BackupDataInput backupDataInput = new BackupDataInput(backupData.getFileDescriptor());
+            while (backupDataInput.readNextHeader()) {
+                String key = backupDataInput.getKey();
+                if (key != null && key.charAt(0) >= 0xff00) {
+                    mReporter.onAgentIllegalKey(mCurrentPackage, key);
+                    // Crash them if they wrote any protected keys.
+                    agentFail(mAgentBinder, "Illegal backup key: " + key);
+                    return false;
+                }
+                backupDataInput.skipEntityData();
+            }
+        }
+        return true;
+    }
+
+    private int getPerformBackupFlags(boolean userInitiated, boolean nonIncremental) {
+        int userInitiatedFlag = userInitiated ? BackupTransport.FLAG_USER_INITIATED : 0;
+        int incrementalFlag =
+                nonIncremental
+                        ? BackupTransport.FLAG_NON_INCREMENTAL
+                        : BackupTransport.FLAG_INCREMENTAL;
+        return userInitiatedFlag | incrementalFlag;
     }
 
     /**
@@ -1023,12 +1027,12 @@
         errorCleanup();
     }
 
-    private void revertBackup() {
-        mReporter.onRevertBackup();
+    private void revertTask() {
+        mReporter.onRevertTask();
         long delay;
         try {
             IBackupTransport transport =
-                    mTransportClient.connectOrThrow("KVBT.revertBackup()");
+                    mTransportClient.connectOrThrow("KVBT.revertTask()");
             delay = transport.requestBackupTime();
         } catch (Exception e) {
             mReporter.onTransportRequestBackupTimeError(e);
diff --git a/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupReporterTest.java b/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupReporterTest.java
new file mode 100644
index 0000000..72ba439
--- /dev/null
+++ b/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupReporterTest.java
@@ -0,0 +1,80 @@
+/*
+ * Copyright (C) 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+
+package com.android.server.backup.keyvalue;
+
+import static com.android.server.backup.keyvalue.KeyValueBackupReporter.TAG;
+import static com.android.server.backup.testing.TestUtils.assertLogcat;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.app.backup.IBackupManagerMonitor;
+import android.app.backup.IBackupObserver;
+import android.platform.test.annotations.Presubmit;
+import android.util.Log;
+
+import com.android.server.backup.BackupManagerService;
+import com.android.server.testing.FrameworkRobolectricTestRunner;
+import com.android.server.testing.SystemLoaderPackages;
+import com.android.server.testing.shadows.ShadowEventLog;
+import com.android.server.testing.shadows.ShadowSlog;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.robolectric.annotation.Config;
+
+@RunWith(FrameworkRobolectricTestRunner.class)
+@Config(
+        manifest = Config.NONE,
+        sdk = 26,
+        shadows = {ShadowEventLog.class, ShadowSlog.class})
+@SystemLoaderPackages({"com.android.server.backup"})
+@Presubmit
+public class KeyValueBackupReporterTest {
+    @Mock private BackupManagerService mBackupManagerService;
+    @Mock private IBackupObserver mObserver;
+    @Mock private IBackupManagerMonitor mMonitor;
+
+    private KeyValueBackupReporter mReporter;
+
+    @Before
+    public void setUp() throws Exception {
+        mReporter = new KeyValueBackupReporter(mBackupManagerService, mObserver, mMonitor);
+    }
+
+    @Test
+    public void testOnNewThread_logsCorrectly() throws Exception {
+        KeyValueBackupReporter.onNewThread("foo");
+
+        assertLogcat(TAG, Log.DEBUG);
+    }
+
+    @Test
+    public void testGetMonitor_returnsMonitor() throws Exception {
+        IBackupManagerMonitor monitor = mReporter.getMonitor();
+
+        assertThat(monitor).isEqualTo(mMonitor);
+    }
+
+    @Test
+    public void testGetObserver_returnsObserver() throws Exception {
+        IBackupObserver observer = mReporter.getObserver();
+
+        assertThat(observer).isEqualTo(mObserver);
+    }
+}
diff --git a/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java b/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
index 9d6b8d5..63b0ea8 100644
--- a/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
+++ b/services/robotests/src/com/android/server/backup/keyvalue/KeyValueBackupTaskTest.java
@@ -36,7 +36,9 @@
 import static com.android.server.backup.testing.TestUtils.uncheck;
 import static com.android.server.backup.testing.TestUtils.waitUntil;
 import static com.android.server.backup.testing.TransportData.backupTransport;
+import static com.android.server.backup.testing.Utils.isFileNonEmpty;
 import static com.android.server.backup.testing.Utils.oneTimeIterable;
+import static com.android.server.backup.testing.Utils.transferStreamedData;
 
 import static com.google.common.truth.Truth.assertThat;
 
@@ -109,8 +111,6 @@
 import com.android.server.backup.testing.TransportData;
 import com.android.server.backup.testing.TransportTestUtils;
 import com.android.server.backup.testing.TransportTestUtils.TransportMock;
-import com.android.server.backup.testing.Utils;
-import com.android.server.backup.transport.TransportClient;
 import com.android.server.testing.FrameworkRobolectricTestRunner;
 import com.android.server.testing.SystemLoaderClasses;
 import com.android.server.testing.SystemLoaderPackages;
@@ -121,6 +121,7 @@
 
 import com.google.common.truth.IterableSubject;
 
+import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -184,6 +185,7 @@
     private ShadowLooper mShadowBackupLooper;
     private Handler mBackupHandler;
     private PowerManager.WakeLock mWakeLock;
+    private KeyValueBackupReporter mReporter;
     private ShadowPackageManager mShadowPackageManager;
     private FakeIBackupManager mBackupManager;
     private File mBaseStateDir;
@@ -220,7 +222,6 @@
         mShadowPackageManager = shadowOf(packageManager);
 
         mWakeLock = createBackupWakeLock(mApplication);
-
         mBackupManager = spy(FakeIBackupManager.class);
 
         // Needed to be able to use a real BMS instead of a mock
@@ -244,15 +245,19 @@
         mBackupHandler = mBackupManagerService.getBackupHandler();
         mShadowBackupLooper = shadowOf(mBackupHandler.getLooper());
         ShadowEventLog.setUp();
+        mReporter = spy(new KeyValueBackupReporter(mBackupManagerService, mObserver, mMonitor));
+    }
+
+    @After
+    public void tearDown() throws Exception {
+        ShadowBackupDataInput.reset();
     }
 
     @Test
     public void testRunTask_whenQueueEmpty_updatesBookkeeping() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         when(mBackupManagerService.getCurrentToken()).thenReturn(0L);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, true);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true);
 
         runTask(task);
 
@@ -266,9 +271,7 @@
     public void testRunTask_whenQueueEmpty_releasesWakeLock() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         when(mBackupManagerService.getCurrentToken()).thenReturn(0L);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, true);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true);
 
         runTask(task);
 
@@ -279,9 +282,7 @@
     public void testRunTask_whenQueueEmpty_doesNotProduceData() throws Exception {
         TransportMock transportMock = setUpTransport(mTransport);
         when(mBackupManagerService.getCurrentToken()).thenReturn(0L);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, true);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true);
 
         runTask(task);
 
@@ -293,9 +294,7 @@
     public void testRunTask_whenQueueEmpty_doesNotCallTransport() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         when(mBackupManagerService.getCurrentToken()).thenReturn(0L);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, true);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true);
 
         runTask(task);
 
@@ -308,9 +307,7 @@
     public void testRunTask_whenQueueEmpty_notifiesCorrectly() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         when(mBackupManagerService.getCurrentToken()).thenReturn(0L);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, true);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true);
 
         runTask(task);
 
@@ -322,9 +319,7 @@
     @Test
     public void testRunTask_whenQueueEmpty_doesNotChangeStateFiles() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, true);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true);
         Files.write(getStateFile(mTransport, PM_PACKAGE), "pmState".getBytes());
         Files.write(getStateFile(mTransport, PACKAGE_1), "packageState".getBytes());
 
@@ -340,9 +335,7 @@
     public void testRunTask_whenOnePackageAndTransportUnavailable() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport.unavailable());
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -355,9 +348,7 @@
     public void testRunTask_whenOnePackage_logsBackupStartEvent() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -368,9 +359,7 @@
     public void testRunTask_whenOnePackage_releasesWakeLock() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -384,9 +373,7 @@
         mBackupManagerService.setCurrentToken(0L);
         when(transportMock.transport.getCurrentRestoreSet()).thenReturn(1234L);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -403,12 +390,7 @@
             throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        false,
-                        PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, false, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
         runTask(task);
@@ -421,12 +403,7 @@
             throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        true,
-                        PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
         runTask(task);
@@ -440,12 +417,7 @@
         setUpAgentWithData(PACKAGE_1);
         PackageManagerBackupAgent pmAgent = spy(createPmAgent());
         when(mBackupManagerService.makeMetadataAgent()).thenReturn(forward(pmAgent));
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        true,
-                        PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true, PACKAGE_1);
 
         runTask(task);
 
@@ -459,12 +431,7 @@
         PackageManagerBackupAgent pmAgent = spy(createPmAgent());
         when(mBackupManagerService.makeMetadataAgent()).thenReturn(forward(pmAgent));
         KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        true,
-                        PACKAGE_1,
-                        PM_PACKAGE);
+                createKeyValueBackupTask(transportMock, true, PACKAGE_1, PM_PACKAGE);
 
         runTask(task);
 
@@ -477,12 +444,7 @@
         setUpAgentWithData(PACKAGE_1);
         PackageManagerBackupAgent pmAgent = spy(createPmAgent());
         when(mBackupManagerService.makeMetadataAgent()).thenReturn(forward(pmAgent));
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        false,
-                        PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, false, PACKAGE_1);
 
         runTask(task);
 
@@ -495,9 +457,7 @@
         TransportMock transportMock = setUpTransport(mTransport);
         // Need 2 packages to be able to verify state of package not involved in the task
         setUpAgentsWithData(PACKAGE_1, PACKAGE_2);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         deletePmStateFile();
         Files.write(getStateFile(mTransport, PACKAGE_2), "package2State".getBytes());
 
@@ -516,9 +476,7 @@
             throws Exception {
         TransportMock transportMock = setUpTransport(mTransport);
         setUpAgentsWithData(PACKAGE_1, PACKAGE_2);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         createPmStateFile();
         Files.write(getStateFile(mTransport, PACKAGE_2), "package2State".getBytes());
 
@@ -535,9 +493,7 @@
         when(transportMock.transport.initializeDevice())
                 .thenReturn(BackupTransport.TRANSPORT_ERROR);
         AgentMock agentMock = setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         deletePmStateFile();
 
         runTask(task);
@@ -556,9 +512,7 @@
         TransportMock transportMock = setUpTransport(mTransport);
         when(transportMock.transport.initializeDevice()).thenThrow(RemoteException.class);
         AgentMock agentMock = setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         deletePmStateFile();
 
         runTask(task);
@@ -575,9 +529,7 @@
     public void testRunTask_whenPackageNotEligibleForBackup() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgentWithData(PACKAGE_1.backupNotAllowed());
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -594,9 +546,7 @@
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         PackageData packageData = fullBackupPackage(1);
         AgentMock agentMock = setUpAgentWithData(packageData);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, packageData);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, packageData);
 
         runTask(task);
 
@@ -611,9 +561,7 @@
     public void testRunTask_whenPackageIsStopped() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgentWithData(PACKAGE_1.stopped());
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -627,9 +575,7 @@
     public void testRunTask_whenPackageUnknown() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         // Not calling setUpAgent()
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -655,9 +601,7 @@
                                     argThat(workSource -> workSource.get(0) == PACKAGE_1.uid));
                     verify(mBackupManagerService, never()).setWorkSource(null);
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -675,9 +619,7 @@
     public void testRunTask_whenAgentUnavailable() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgent(PACKAGE_1.unavailable());
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -693,9 +635,7 @@
         doThrow(SecurityException.class)
                 .when(mBackupManagerService)
                 .bindToAgentSynchronous(argThat(applicationInfo(PACKAGE_1)), anyInt());
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -710,9 +650,7 @@
         when(transportMock.transport.getBackupQuota(PACKAGE_1.packageName, false))
                 .thenThrow(DeadObjectException.class);
         setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -727,9 +665,7 @@
         when(transportMock.transport.getBackupQuota(PACKAGE_1.packageName, false))
                 .thenThrow(DeadObjectException.class);
         setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -743,9 +679,7 @@
         when(transportMock.transport.getBackupQuota(PACKAGE_1.packageName, false))
                 .thenThrow(DeadObjectException.class);
         setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "packageState".getBytes());
 
         runTask(task);
@@ -757,20 +691,16 @@
     }
 
     @Test
-    public void testRunTask_whenTransportGetBackupQuotaThrows_revertsOperation() throws Exception {
+    public void testRunTask_whenTransportGetBackupQuotaThrows_revertsTask() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         when(transportMock.transport.getBackupQuota(PACKAGE_1.packageName, false))
                 .thenThrow(DeadObjectException.class);
         setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
-        verify(transportMock.transport).requestBackupTime();
-        assertBackupPendingFor(PACKAGE_1);
-        assertThat(KeyValueBackupJob.isScheduled()).isTrue();
+        assertTaskReverted(transportMock, PACKAGE_1);
     }
 
     /**
@@ -789,9 +719,7 @@
                 (oldState, dataOutput, newState) -> {
                     throw new RuntimeException();
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -811,9 +739,7 @@
         int flags = BackupAgent.FLAG_CLIENT_SIDE_ENCRYPTION_ENABLED;
         when(transportMock.transport.getTransportFlags()).thenReturn(flags);
         AgentMock agentMock = setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -825,9 +751,7 @@
     public void testRunTask_whenTransportDoesNotProvidesFlags() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -843,12 +767,7 @@
         List<AgentMock> agentMocks = setUpAgents(PACKAGE_1, PACKAGE_2);
         BackupAgent agent1 = agentMocks.get(0).agent;
         BackupAgent agent2 = agentMocks.get(1).agent;
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        PACKAGE_1,
-                        PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
 
         runTask(task);
 
@@ -860,9 +779,7 @@
     public void testRunTask_whenTransportChangeFlagsAfterTaskCreation() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         int flags = BackupAgent.FLAG_CLIENT_SIDE_ENCRYPTION_ENABLED;
         when(transportMock.transport.getTransportFlags()).thenReturn(flags);
 
@@ -883,9 +800,7 @@
                     writeData(dataOutput, prohibitedChar + "key", "data".getBytes());
                     writeState(newState, "newState".getBytes());
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -904,9 +819,7 @@
                     writeData(dataOutput, prohibitedChar + "key", "data".getBytes());
                     writeState(newState, "newState".getBytes());
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
         runTask(task);
@@ -928,9 +841,7 @@
                     writeData(dataOutput, prohibitedChar + "key", "data".getBytes());
                     writeState(newState, "newState".getBytes());
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
         runTask(task);
@@ -950,9 +861,7 @@
                     writeData(dataOutput, prohibitedChar + "key", "data".getBytes());
                     writeState(newState, "newState".getBytes());
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
         runTask(task);
@@ -973,9 +882,7 @@
                     writeData(dataOutput, prohibitedChar + "key", "data".getBytes());
                     writeState(newState, "newState".getBytes());
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
         runTask(task);
@@ -1003,12 +910,7 @@
                     writeData(dataOutput, "key", "data".getBytes());
                     writeState(newState, "newState".getBytes());
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        PACKAGE_1,
-                        PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
 
         runTask(task);
 
@@ -1028,9 +930,7 @@
                 (oldState, dataOutput, newState) -> {
                     // No-op
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1047,9 +947,7 @@
                 (oldState, dataOutput, newState) -> {
                     // No-op
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1066,9 +964,7 @@
                 (oldState, dataOutput, newState) -> {
                     // No-op
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1085,9 +981,7 @@
                 (oldState, dataOutput, newState) -> {
                     // No-op
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1103,9 +997,7 @@
                 (oldState, dataOutput, newState) -> {
                     // No-op
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1130,9 +1022,7 @@
                     writeData(dataOutput, "key2", "data2".getBytes());
                     writeState(newState, "newState".getBytes());
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1155,9 +1045,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_OK);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1176,9 +1064,7 @@
                     writeData(dataOutput, "key", "data".getBytes());
                     writeState(newState, "newState".getBytes());
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1196,9 +1082,7 @@
         when(transportMock.transport.performBackup(
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .then(copyBackupDataTo(backupData));
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1210,9 +1094,7 @@
     public void testRunTask_whenFinishBackupSucceeds_notifiesCorrectly() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1226,9 +1108,7 @@
     public void testRunTask_whenFinishBackupSucceeds_updatesBookkeeping() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1242,9 +1122,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_PACKAGE_REJECTED);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1259,9 +1137,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_PACKAGE_REJECTED);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
         runTask(task);
@@ -1279,9 +1155,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_PACKAGE_REJECTED);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1296,9 +1170,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_PACKAGE_REJECTED);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1314,9 +1186,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_PACKAGE_REJECTED);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1333,12 +1203,7 @@
                         argThat(packageInfo(PACKAGE_2)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_OK);
         setUpAgentsWithData(PACKAGE_1, PACKAGE_2);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        PACKAGE_1,
-                        PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
 
         runTask(task);
 
@@ -1357,12 +1222,7 @@
                         argThat(packageInfo(PACKAGE_2)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_PACKAGE_REJECTED);
         setUpAgentsWithData(PACKAGE_1, PACKAGE_2);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        PACKAGE_1,
-                        PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
 
         runTask(task);
 
@@ -1372,7 +1232,8 @@
     }
 
     @Test
-    public void testRunTask_whenTransportReturnsQuotaExceeded() throws Exception {
+    public void testRunTask_whenTransportReturnsQuotaExceeded_callsAgentOnQuotaExceeded()
+            throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         when(transportMock.transport.getBackupQuota(PACKAGE_1.packageName, false))
                 .thenReturn(1234L);
@@ -1380,19 +1241,206 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_QUOTA_EXCEEDED);
         AgentMock agentMock = setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        verify(agentMock.agent).onQuotaExceeded(anyLong(), eq(1234L));
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsQuotaExceeded_updatesBookkeeping()
+            throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_QUOTA_EXCEEDED);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        assertBackupNotPendingFor(PACKAGE_1);
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsQuotaExceeded_notifiesAndLogs() throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.getBackupQuota(PACKAGE_1.packageName, false))
+                .thenReturn(1234L);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_QUOTA_EXCEEDED);
+        AgentMock agentMock = setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
         verify(mObserver)
                 .onResult(PACKAGE_1.packageName, BackupManager.ERROR_TRANSPORT_QUOTA_EXCEEDED);
         verify(mObserver).backupFinished(SUCCESS);
-        verify(agentMock.agent).onQuotaExceeded(anyLong(), eq(1234L));
         assertEventLogged(EventLogTags.BACKUP_QUOTA_EXCEEDED, PACKAGE_1.packageName);
-        assertBackupNotPendingFor(PACKAGE_1);
-        // TODO: Assert about state/staging files (possible bug)
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsQuotaExceeded_cleansUpFiles() throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_QUOTA_EXCEEDED);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        assertThat(Files.exists(getTemporaryStateFile(mTransport, PACKAGE_1))).isFalse();
+        assertThat(Files.exists(getStagingFile(PACKAGE_1))).isFalse();
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsNotInitialized_cleansUpFiles() throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        assertThat(isFileNonEmpty(getStateFile(mTransport, PACKAGE_1))).isFalse();
+        assertThat(Files.exists(getTemporaryStateFile(mTransport, PACKAGE_1))).isFalse();
+        assertThat(Files.exists(getStagingFile(PACKAGE_1))).isFalse();
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsNotInitialized_reportsCorrectly() throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        verify(mReporter).onPackageBackupTransportFailure(PACKAGE_1.packageName);
+        verify(mReporter).onTransportNotInitialized();
+        verify(mReporter).onBackupFinished(BackupManager.ERROR_TRANSPORT_ABORTED);
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsNotInitializedForPm_reportsCorrectly()
+            throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PM_PACKAGE)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        verify(mReporter).onPackageBackupTransportFailure(PM_PACKAGE.packageName);
+        verify(mReporter).onTransportNotInitialized();
+        verify(mReporter).onBackupFinished(BackupManager.ERROR_TRANSPORT_ABORTED);
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsNotInitialized_doesNotCallSecondAgent()
+            throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        setUpAgentWithData(PACKAGE_1);
+        AgentMock agentMock = setUpAgentWithData(PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
+
+        runTask(task);
+
+        verify(agentMock.agent, never()).onBackup(any(), any(), any());
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsNotInitialized_revertsTask() throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        assertTaskReverted(transportMock, PACKAGE_1);
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsNotInitialized_triggersTransportInitialization()
+            throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        assertThat(mBackupManagerService.getPendingInits()).contains(mTransport.transportName);
+        verify(mBackupManagerService).backupNow();
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsNotInitialized_cleansUpPmStateFile()
+            throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+        Files.write(getStateFile(mTransport, PM_PACKAGE), "pmState".getBytes());
+
+        runTask(task);
+
+        assertThat(Files.exists(getStateFile(mTransport, PM_PACKAGE))).isFalse();
+    }
+
+    @Test
+    public void testRunTask_whenTransportReturnsNotInitializedForPm_cleansUpPmStateFile()
+            throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(
+                        argThat(packageInfo(PM_PACKAGE)), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+        Files.write(getStateFile(mTransport, PM_PACKAGE), "pmState".getBytes());
+
+        runTask(task);
+
+        assertThat(Files.exists(getStateFile(mTransport, PM_PACKAGE))).isFalse();
+    }
+
+    @Test
+    public void
+            testRunTask_whenTransportReturnsNotInitializedAndThrowsWhenQueryingName_reportsCorrectly()
+                    throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        when(transportMock.transport.performBackup(any(), any(), anyInt()))
+                .thenReturn(BackupTransport.TRANSPORT_NOT_INITIALIZED);
+        // First one is in startTask(), second is the one we want.
+        when(transportMock.transport.name())
+                .thenReturn(mTransport.transportName)
+                .thenThrow(DeadObjectException.class);
+        setUpAgentWithData(PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
+
+        runTask(task);
+
+        verify(mReporter).onPendingInitializeTransportError(any(DeadObjectException.class));
+        verify(mReporter).onBackupFinished(ERROR_TRANSPORT_ABORTED);
     }
 
     @Test
@@ -1403,12 +1451,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_NON_INCREMENTAL_BACKUP_REQUIRED);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        true,
-                        PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, true, PACKAGE_1);
         // Delete to be non-incremental
         Files.deleteIfExists(getStateFile(mTransport, PACKAGE_1));
 
@@ -1451,12 +1494,7 @@
                         writeState(newState, "stateForNonIncremental".getBytes());
                     }
                 });
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        false,
-                        PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, false, PACKAGE_1);
         // Write state to be incremental
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
@@ -1490,9 +1528,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_ERROR);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1508,9 +1544,7 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_ERROR);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1518,21 +1552,17 @@
     }
 
     @Test
-    public void testRunTask_whenTransportReturnsError_revertsOperation() throws Exception {
+    public void testRunTask_whenTransportReturnsError_revertsTask() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         when(transportMock.transport.performBackup(
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_ERROR);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
-        verify(transportMock.transport).requestBackupTime();
-        assertBackupPendingFor(PACKAGE_1);
-        assertThat(KeyValueBackupJob.isScheduled()).isTrue();
+        assertTaskReverted(transportMock, PACKAGE_1);
     }
 
     @Test
@@ -1542,18 +1572,15 @@
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenReturn(BackupTransport.TRANSPORT_ERROR);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         Files.write(getStateFile(mTransport, PACKAGE_1), "oldState".getBytes());
 
         runTask(task);
 
         assertThat(Files.readAllBytes(getStateFile(mTransport, PACKAGE_1)))
                 .isEqualTo("oldState".getBytes());
-        // TODO: These should be true (Bug)
-        // assertThat(Files.exists(getTemporaryStateFile(mTransport, PACKAGE_1))).isFalse();
-        // assertThat(Files.exists(getStagingFile(PACKAGE_1))).isFalse();
+        assertThat(Files.exists(getTemporaryStateFile(mTransport, PACKAGE_1))).isFalse();
+        assertThat(Files.exists(getStagingFile(PACKAGE_1))).isFalse();
     }
 
     @Test
@@ -1562,9 +1589,7 @@
         when(transportMock.transport.getBackupQuota(PM_PACKAGE.packageName, false))
                 .thenThrow(DeadObjectException.class);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1581,9 +1606,7 @@
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         PackageManagerBackupAgent pmAgent = createThrowingPmAgent(new RuntimeException());
         when(mBackupManagerService.makeMetadataAgent()).thenReturn(pmAgent);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
 
         runTask(task);
 
@@ -1599,26 +1622,37 @@
     public void testRunTask_whenBackupRunning_doesNotThrow() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         when(mBackupManagerService.isBackupOperationInProgress()).thenReturn(true);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock);
 
         runTask(task);
     }
 
     @Test
+    public void testRunTask_whenReadingBackupDataThrows() throws Exception {
+        TransportMock transportMock = setUpInitializedTransport(mTransport);
+        setUpAgentWithData(PACKAGE_1);
+        AgentMock agentMock = setUpAgentWithData(PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
+        // We don't validate PM's data, so it will only throw in PACKAGE_1
+        ShadowBackupDataInput.throwInNextHeaderRead();
+
+        runTask(task);
+
+        verify(mReporter).onReadAgentDataError(eq(PACKAGE_1.packageName), any());
+        verify(transportMock.transport, never())
+                .performBackup(argThat(packageInfo(PACKAGE_1)), any(), anyInt());
+        verify(agentMock.agent, never()).onBackup(any(), any(), any());
+        assertTaskReverted(transportMock, PACKAGE_1, PACKAGE_2);
+    }
+
+    @Test
     public void
             testRunTask_whenMarkCancelDuringFirstAgentOnBackup_doesNotCallTransportAfterWaitCancel()
                     throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgent(PACKAGE_1);
         setUpAgentsWithData(PACKAGE_2);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        PACKAGE_1,
-                        PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
         agentOnBackupDo(
                 agentMock,
                 (oldState, dataOutput, newState) -> {
@@ -1641,9 +1675,7 @@
             throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         agentOnBackupDo(
                 agentMock,
                 (oldState, dataOutput, newState) -> {
@@ -1667,9 +1699,7 @@
             throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         when(transportMock.transport.performBackup(
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenAnswer(
@@ -1695,12 +1725,7 @@
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
         AgentMock agentMock = setUpAgent(PACKAGE_2);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        PACKAGE_1,
-                        PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
         agentOnBackupDo(
                 agentMock,
                 (oldState, dataOutput, newState) -> {
@@ -1727,12 +1752,7 @@
                     throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentsWithData(PACKAGE_1, PACKAGE_2);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient,
-                        mTransport.transportDirName,
-                        PACKAGE_1,
-                        PACKAGE_2);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1, PACKAGE_2);
         when(transportMock.transport.performBackup(
                         argThat(packageInfo(PACKAGE_1)), any(), anyInt()))
                 .thenAnswer(
@@ -1757,9 +1777,7 @@
     public void testRunTask_afterMarkCancel_doesNotCallAgentOrTransport() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         AgentMock agentMock = setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         task.markCancel();
 
         runTask(task);
@@ -1773,9 +1791,7 @@
     public void testWaitCancel_afterCancelledTaskFinished_returns() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         task.markCancel();
         runTask(task);
 
@@ -1783,13 +1799,12 @@
     }
 
     @Test
-    public void testWaitCancel_whenMarkCancelDuringAgentOnBackup_unregistersTask() throws Exception {
+    public void testWaitCancel_whenMarkCancelDuringAgentOnBackup_unregistersTask()
+            throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
         AgentMock agentMock = setUpAgent(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         agentOnBackupDo(
                 agentMock,
                 (oldState, dataOutput, newState) -> {
@@ -1816,9 +1831,7 @@
     public void testMarkCancel_afterTaskFinished_returns() throws Exception {
         TransportMock transportMock = setUpInitializedTransport(mTransport);
         setUpAgentWithData(PACKAGE_1);
-        KeyValueBackupTask task =
-                createKeyValueBackupTask(
-                        transportMock.transportClient, mTransport.transportDirName, PACKAGE_1);
+        KeyValueBackupTask task = createKeyValueBackupTask(transportMock, PACKAGE_1);
         runTask(task);
 
         task.markCancel();
@@ -1997,15 +2010,12 @@
     }
 
     private KeyValueBackupTask createKeyValueBackupTask(
-            TransportClient transportClient, String transportDirName, PackageData... packages) {
-        return createKeyValueBackupTask(transportClient, transportDirName, false, packages);
+            TransportMock transportMock, PackageData... packages) {
+        return createKeyValueBackupTask(transportMock, false, packages);
     }
 
     private KeyValueBackupTask createKeyValueBackupTask(
-            TransportClient transportClient,
-            String transportDirName,
-            boolean nonIncremental,
-            PackageData... packages) {
+            TransportMock transportMock, boolean nonIncremental, PackageData... packages) {
         List<String> queue =
                 Stream.of(packages).map(packageData -> packageData.packageName).collect(toList());
         mBackupManagerService.getPendingBackups().clear();
@@ -2015,12 +2025,11 @@
         KeyValueBackupTask task =
                 new KeyValueBackupTask(
                         mBackupManagerService,
-                        transportClient,
-                        transportDirName,
+                        transportMock.transportClient,
+                        transportMock.transportData.transportDirName,
                         queue,
                         mOldJournal,
-                        mObserver,
-                        mMonitor,
+                        mReporter,
                         mListener,
                         emptyList(),
                         /* userInitiated */ false,
@@ -2129,7 +2138,7 @@
                                 (oldState, dataOutput, newState) -> {
                                     ByteArrayOutputStream outputStream =
                                             new ByteArrayOutputStream();
-                                    Utils.transferStreamedData(
+                                    transferStreamedData(
                                             new FileInputStream(oldState.getFileDescriptor()),
                                             outputStream);
                                     agentMock.oldState = outputStream.toByteArray();
@@ -2180,18 +2189,31 @@
         assertThat(packages).doesNotContain(packageName);
     }
 
-    private void assertBackupPendingFor(PackageData packageData) throws IOException {
-        String packageName = packageData.packageName;
-        // We verify the current journal, NOT the old one passed to KeyValueBackupTask constructor
-        assertThat(mBackupManagerService.getJournal().getPackages()).contains(packageName);
-        assertThat(mBackupManagerService.getPendingBackups()).containsKey(packageName);
+    private void assertTaskReverted(TransportMock transportMock, PackageData... packages)
+            throws RemoteException, IOException {
+        verify(transportMock.transport).requestBackupTime();
+        assertBackupPendingFor(packages);
+        assertThat(KeyValueBackupJob.isScheduled()).isTrue();
     }
 
-    private void assertBackupNotPendingFor(PackageData packageData) throws IOException {
-        String packageName = packageData.packageName;
-        // We verify the current journal, NOT the old one passed to KeyValueBackupTask constructor
-        assertJournalDoesNotContain(mBackupManagerService.getJournal(), packageName);
-        assertThat(mBackupManagerService.getPendingBackups()).doesNotContainKey(packageName);
+    private void assertBackupPendingFor(PackageData... packages) throws IOException {
+        for (PackageData packageData : packages) {
+            String packageName = packageData.packageName;
+            // We verify the current journal, NOT the old one passed to KeyValueBackupTask
+            // constructor
+            assertThat(mBackupManagerService.getJournal().getPackages()).contains(packageName);
+            assertThat(mBackupManagerService.getPendingBackups()).containsKey(packageName);
+        }
+    }
+
+    private void assertBackupNotPendingFor(PackageData... packages) throws IOException {
+        for (PackageData packageData : packages) {
+            String packageName = packageData.packageName;
+            // We verify the current journal, NOT the old one passed to KeyValueBackupTask
+            // constructor
+            assertJournalDoesNotContain(mBackupManagerService.getJournal(), packageName);
+            assertThat(mBackupManagerService.getPendingBackups()).doesNotContainKey(packageName);
+        }
     }
 
     private void assertDataHasKeyValue(BackupDataInput backupData, String key, byte[] value)
diff --git a/services/robotests/src/com/android/server/backup/testing/TestUtils.java b/services/robotests/src/com/android/server/backup/testing/TestUtils.java
index df4d457..2f54513 100644
--- a/services/robotests/src/com/android/server/backup/testing/TestUtils.java
+++ b/services/robotests/src/com/android/server/backup/testing/TestUtils.java
@@ -20,6 +20,8 @@
 
 import static org.robolectric.Shadows.shadowOf;
 
+import static java.util.stream.Collectors.toSet;
+
 import android.os.Looper;
 import android.os.Message;
 import android.os.MessageQueue;
@@ -36,6 +38,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
+import java.util.stream.IntStream;
 
 public class TestUtils {
     private static final long TIMEOUT_MS = 3000;
@@ -87,20 +90,45 @@
         ShadowSystemClock.setCurrentTimeMillis(shadowLooper.getScheduler().getCurrentTime());
     }
 
-    /** Reset logcat with {@link ShadowLog#reset()} before the test case. */
+    /**
+     * Reset logcat with {@link ShadowLog#reset()} before the test case if you do anything that uses
+     * logcat before that.
+     */
     public static void assertLogcatAtMost(String tag, int level) {
         assertThat(ShadowLog.getLogsForTag(tag).stream().allMatch(logItem -> logItem.type <= level))
                 .named("All logs <= " + level)
                 .isTrue();
     }
 
-    /** Reset logcat with {@link ShadowLog#reset()} before the test case. */
+    /**
+     * Reset logcat with {@link ShadowLog#reset()} before the test case if you do anything that uses
+     * logcat before that.
+     */
     public static void assertLogcatAtLeast(String tag, int level) {
         assertThat(ShadowLog.getLogsForTag(tag).stream().anyMatch(logItem -> logItem.type >= level))
                 .named("Any log >= " + level)
                 .isTrue();
     }
 
+    /**
+     * Verifies that logcat has produced log items as specified per level in {@code logs} (with
+     * repetition).
+     *
+     * <p>So, if you call {@code assertLogcat(TAG, Log.ERROR, Log.ERROR)}, you assert that there are
+     * exactly 2 log items, each with level ERROR.
+     *
+     * <p>Reset logcat with {@link ShadowLog#reset()} before the test case if you do anything
+     * that uses logcat before that.
+     */
+    public static void assertLogcat(String tag, int... logs) {
+        assertThat(
+                        ShadowLog.getLogsForTag(tag).stream()
+                                .map(logItem -> logItem.type)
+                                .collect(toSet()))
+                .named("Log items (specified per level)")
+                .containsExactly(IntStream.of(logs).boxed().toArray());
+    }
+
     public static void assertLogcatContains(String tag, Predicate<ShadowLog.LogItem> predicate) {
         assertThat(ShadowLog.getLogsForTag(tag).stream().anyMatch(predicate)).isTrue();
     }
diff --git a/services/robotests/src/com/android/server/backup/testing/TransportTestUtils.java b/services/robotests/src/com/android/server/backup/testing/TransportTestUtils.java
index 6625443..f6ed630 100644
--- a/services/robotests/src/com/android/server/backup/testing/TransportTestUtils.java
+++ b/services/robotests/src/com/android/server/backup/testing/TransportTestUtils.java
@@ -164,18 +164,18 @@
                 when(transportClientMock.connectOrThrow(any())).thenReturn(transportMock);
                 when(transportClientMock.connect(any())).thenReturn(transportMock);
 
-                return new TransportMock(transportClientMock, transportMock);
+                return new TransportMock(transport, transportClientMock, transportMock);
             } else {
                 // Transport registered but unavailable
                 when(transportClientMock.connectOrThrow(any()))
                         .thenThrow(TransportNotAvailableException.class);
                 when(transportClientMock.connect(any())).thenReturn(null);
 
-                return new TransportMock(transportClientMock, null);
+                return new TransportMock(transport, transportClientMock, null);
             }
         } else {
             // Transport not registered
-            return new TransportMock(null, null);
+            return new TransportMock(transport, null, null);
         }
     }
 
@@ -196,11 +196,15 @@
     }
 
     public static class TransportMock {
+        public final TransportData transportData;
         @Nullable public final TransportClient transportClient;
         @Nullable public final IBackupTransport transport;
 
         private TransportMock(
-                @Nullable TransportClient transportClient, @Nullable IBackupTransport transport) {
+                TransportData transportData,
+                @Nullable TransportClient transportClient,
+                @Nullable IBackupTransport transport) {
+            this.transportData = transportData;
             this.transportClient = transportClient;
             this.transport = transport;
         }
diff --git a/services/robotests/src/com/android/server/backup/testing/Utils.java b/services/robotests/src/com/android/server/backup/testing/Utils.java
index bd8b4ef..b0e00a2 100644
--- a/services/robotests/src/com/android/server/backup/testing/Utils.java
+++ b/services/robotests/src/com/android/server/backup/testing/Utils.java
@@ -19,6 +19,8 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.Iterator;
 
 public class Utils {
@@ -41,5 +43,9 @@
         return () -> iterator;
     }
 
+    public static boolean isFileNonEmpty(Path path) throws IOException {
+        return Files.exists(path) && Files.size(path) > 0;
+    }
+
     private Utils() {}
 }
diff --git a/services/robotests/src/com/android/server/testing/shadows/ShadowBackupDataInput.java b/services/robotests/src/com/android/server/testing/shadows/ShadowBackupDataInput.java
index bc47dd5..4901828 100644
--- a/services/robotests/src/com/android/server/testing/shadows/ShadowBackupDataInput.java
+++ b/services/robotests/src/com/android/server/testing/shadows/ShadowBackupDataInput.java
@@ -20,6 +20,7 @@
 
 import org.robolectric.annotation.Implementation;
 import org.robolectric.annotation.Implements;
+import org.robolectric.annotation.Resetter;
 
 import java.io.EOFException;
 import java.io.FileDescriptor;
@@ -33,6 +34,17 @@
  */
 @Implements(BackupDataInput.class)
 public class ShadowBackupDataInput {
+    private static boolean sReadNextHeaderThrow = false;
+
+    public static void throwInNextHeaderRead() {
+        sReadNextHeaderThrow = true;
+    }
+
+    @Resetter
+    public static void reset() {
+        sReadNextHeaderThrow = false;
+    }
+
     private FileDescriptor mFileDescriptor;
     private ObjectInputStream mInput;
     private int mSize;
@@ -46,6 +58,10 @@
 
     @Implementation
     public boolean readNextHeader() throws IOException {
+        if (sReadNextHeaderThrow) {
+            sReadNextHeaderThrow = false;
+            throw new IOException("Fake exception");
+        }
         mHeaderReady = false;
         try {
             ensureInput();
diff --git a/services/robotests/src/com/android/server/testing/shadows/ShadowKeyValueBackupTask.java b/services/robotests/src/com/android/server/testing/shadows/ShadowKeyValueBackupTask.java
index aeda2dc..b7db56b 100644
--- a/services/robotests/src/com/android/server/testing/shadows/ShadowKeyValueBackupTask.java
+++ b/services/robotests/src/com/android/server/testing/shadows/ShadowKeyValueBackupTask.java
@@ -17,12 +17,11 @@
 package com.android.server.testing.shadows;
 
 import android.annotation.Nullable;
-import android.app.backup.IBackupManagerMonitor;
-import android.app.backup.IBackupObserver;
 
 import com.android.server.backup.BackupManagerService;
 import com.android.server.backup.DataChangedJournal;
 import com.android.server.backup.internal.OnTaskFinishedListener;
+import com.android.server.backup.keyvalue.KeyValueBackupReporter;
 import com.android.server.backup.keyvalue.KeyValueBackupTask;
 import com.android.server.backup.transport.TransportClient;
 
@@ -57,12 +56,11 @@
     public void __constructor__(
             BackupManagerService backupManagerService,
             TransportClient transportClient,
-            String dirName,
+            String transportDirName,
             List<String> queue,
-            @Nullable DataChangedJournal journal,
-            IBackupObserver observer,
-            IBackupManagerMonitor monitor,
-            @Nullable OnTaskFinishedListener listener,
+            @Nullable DataChangedJournal dataChangedJournal,
+            KeyValueBackupReporter reporter,
+            OnTaskFinishedListener listener,
             List<String> pendingFullBackups,
             boolean userInitiated,
             boolean nonIncremental) {