Add tests for ActiveRestoreSession

Added first set around getAvailableRestoreSets() plus some tidy-up
around setup and unused stubs.

Test: m -j RunFrameworksServicesRoboTests
Change-Id: I9b3bf548251f1262907de96407184bd1822b3429
diff --git a/services/robotests/src/com/android/server/backup/BackupManagerServiceTest.java b/services/robotests/src/com/android/server/backup/BackupManagerServiceTest.java
index df09780..53d734b 100644
--- a/services/robotests/src/com/android/server/backup/BackupManagerServiceTest.java
+++ b/services/robotests/src/com/android/server/backup/BackupManagerServiceTest.java
@@ -16,6 +16,7 @@
 
 package com.android.server.backup;
 
+import static com.android.server.backup.testing.BackupManagerServiceTestUtils.startBackupThread;
 import static com.android.server.backup.testing.TransportData.backupTransport;
 import static com.android.server.backup.testing.TransportData.d2dTransport;
 import static com.android.server.backup.testing.TransportData.localTransport;
@@ -41,9 +42,9 @@
 import android.os.HandlerThread;
 import android.platform.test.annotations.Presubmit;
 import android.provider.Settings;
-import com.android.server.backup.PackageManagerBackupAgent;
-import com.android.server.backup.testing.ShadowAppBackupUtils;
-import com.android.server.backup.testing.ShadowBackupPolicyEnforcer;
+
+import com.android.server.testing.shadows.ShadowAppBackupUtils;
+import com.android.server.testing.shadows.ShadowBackupPolicyEnforcer;
 import com.android.server.backup.testing.TransportData;
 import com.android.server.backup.testing.TransportTestUtils.TransportMock;
 import com.android.server.backup.transport.TransportNotRegisteredException;
@@ -97,10 +98,7 @@
         mTransport = backupTransport();
         mTransportName = mTransport.transportName;
 
-        mBackupThread = new HandlerThread("backup-test");
-        mBackupThread.setUncaughtExceptionHandler(
-                (t, e) -> ShadowLog.e(TAG, "Uncaught exception in test thread " + t.getName(), e));
-        mBackupThread.start();
+        mBackupThread = startBackupThread(this::uncaughtException);
         mShadowBackupLooper = shadowOf(mBackupThread.getLooper());
 
         ContextWrapper context = RuntimeEnvironment.application;
@@ -121,6 +119,13 @@
         ShadowBackupPolicyEnforcer.setMandatoryBackupTransport(null);
     }
 
+    private void uncaughtException(Thread thread, Throwable e) {
+        // Unrelated exceptions are thrown in the backup thread. Until we mock everything properly
+        // we should not fail tests because of this. This is not flakiness, the exceptions thrown
+        // don't interfere with the tests.
+        ShadowLog.e(TAG, "Uncaught exception in test thread " + thread.getName(), e);
+    }
+
     /* Tests for destination string */
 
     @Test
diff --git a/services/robotests/src/com/android/server/backup/PerformBackupTaskTest.java b/services/robotests/src/com/android/server/backup/PerformBackupTaskTest.java
index a5beed0..d0e6658 100644
--- a/services/robotests/src/com/android/server/backup/PerformBackupTaskTest.java
+++ b/services/robotests/src/com/android/server/backup/PerformBackupTaskTest.java
@@ -16,11 +16,13 @@
 
 package com.android.server.backup;
 
+import static com.android.server.backup.testing.BackupManagerServiceTestUtils.createBackupWakeLock;
+import static com.android.server.backup.testing.BackupManagerServiceTestUtils.setUpBackupManagerServiceBasics;
+import static com.android.server.backup.testing.BackupManagerServiceTestUtils.startBackupThreadAndGetLooper;
 import static com.android.server.backup.testing.TransportData.backupTransport;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyLong;
@@ -28,7 +30,6 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doNothing;
-import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
@@ -41,7 +42,6 @@
 import static java.util.stream.Collectors.toList;
 
 import android.app.Application;
-import android.app.IActivityManager;
 import android.app.IBackupAgent;
 import android.app.backup.BackupAgent;
 import android.app.backup.BackupDataInput;
@@ -52,12 +52,10 @@
 import android.app.backup.IBackupManager;
 import android.app.backup.IBackupManagerMonitor;
 import android.app.backup.IBackupObserver;
-import android.content.Context;
 import android.content.pm.ApplicationInfo;
 import android.content.pm.PackageInfo;
 import android.content.pm.PackageManager;
 import android.os.Handler;
-import android.os.HandlerThread;
 import android.os.Looper;
 import android.os.Message;
 import android.os.ParcelFileDescriptor;
@@ -65,7 +63,6 @@
 import android.os.RemoteException;
 import android.platform.test.annotations.Presubmit;
 import android.util.Pair;
-import android.util.SparseArray;
 
 import com.android.internal.backup.IBackupTransport;
 import com.android.server.backup.internal.BackupHandler;
@@ -159,33 +156,24 @@
         PackageManager packageManager = application.getPackageManager();
         mShadowPackageManager = Shadow.extract(packageManager);
 
-        PowerManager powerManager =
-                (PowerManager) application.getSystemService(Context.POWER_SERVICE);
-        mWakeLock = powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, "*backup*");
+        mWakeLock = createBackupWakeLock(application);
 
-        // Robolectric simulates multi-thread in a single-thread to avoid flakiness
-        HandlerThread backupThread = new HandlerThread("backup");
-        backupThread.setUncaughtExceptionHandler(
-                (t, e) -> fail("Uncaught exception " + e.getMessage()));
-        backupThread.start();
-        Looper backupLooper = backupThread.getLooper();
+        Looper backupLooper = startBackupThreadAndGetLooper();
         mShadowBackupLooper = shadowOf(backupLooper);
         mBackupHandler = new BackupHandler(mBackupManagerService, backupLooper);
 
         mBackupManager = spy(FakeIBackupManager.class);
 
-        when(mBackupManagerService.getTransportManager()).thenReturn(mTransportManager);
-        when(mBackupManagerService.getContext()).thenReturn(application);
-        when(mBackupManagerService.getPackageManager()).thenReturn(packageManager);
-        when(mBackupManagerService.getWakelock()).thenReturn(mWakeLock);
-        when(mBackupManagerService.getCurrentOpLock()).thenReturn(new Object());
-        when(mBackupManagerService.getQueueLock()).thenReturn(new Object());
+        setUpBackupManagerServiceBasics(
+                mBackupManagerService,
+                application,
+                mTransportManager,
+                packageManager,
+                mBackupHandler,
+                mWakeLock);
         when(mBackupManagerService.getBaseStateDir()).thenReturn(mBaseStateDir);
         when(mBackupManagerService.getDataDir()).thenReturn(dataDir);
-        when(mBackupManagerService.getCurrentOperations()).thenReturn(new SparseArray<>());
-        when(mBackupManagerService.getBackupHandler()).thenReturn(mBackupHandler);
         when(mBackupManagerService.getBackupManagerBinder()).thenReturn(mBackupManager);
-        when(mBackupManagerService.getActivityManager()).thenReturn(mock(IActivityManager.class));
     }
 
     @Test
@@ -283,6 +271,22 @@
     }
 
     @Test
+    public void testRunTask_releasesWakeLock() throws Exception {
+        TransportMock transportMock = setUpTransport(mTransport);
+        setUpAgent(PACKAGE_1);
+        PerformBackupTask task =
+                createPerformBackupTask(
+                        transportMock.transportClient,
+                        mTransport.transportDirName,
+                        emptyList(),
+                        PACKAGE_1);
+
+        runTask(task);
+
+        assertThat(mWakeLock.isHeld()).isFalse();
+    }
+
+    @Test
     public void testRunTask_callsTransportPerformBackupWithAgentData() throws Exception {
         TransportMock transportMock = setUpTransport(mTransport);
         IBackupTransport transportBinder = transportMock.transport;
diff --git a/services/robotests/src/com/android/server/backup/TransportManagerTest.java b/services/robotests/src/com/android/server/backup/TransportManagerTest.java
index 503adb2..32697bd 100644
--- a/services/robotests/src/com/android/server/backup/TransportManagerTest.java
+++ b/services/robotests/src/com/android/server/backup/TransportManagerTest.java
@@ -49,7 +49,6 @@
 import android.content.pm.PackageInfo;
 import android.platform.test.annotations.Presubmit;
 
-import com.android.server.backup.testing.ShadowContextImplForBackup;
 import com.android.server.backup.testing.TransportData;
 import com.android.server.backup.testing.TransportTestUtils.TransportMock;
 import com.android.server.backup.transport.OnTransportRegisteredListener;
@@ -61,7 +60,6 @@
 import com.android.server.testing.shadows.FrameworkShadowContextImpl;
 import com.android.server.testing.shadows.FrameworkShadowPackageManager;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -90,8 +88,8 @@
     private static final String PACKAGE_B = "some.package.b";
 
     /**
-     * GMSCore depends on this constant so we define it here on top of the definition in
-     * {@link TransportManager} to verify this extra is passed
+     * GMSCore depends on this constant so we define it here on top of the definition in {@link
+     * TransportManager} to verify this extra is passed
      */
     private static final String EXTRA_TRANSPORT_REGISTRATION = "transport_registration";
 
@@ -108,21 +106,15 @@
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
 
+        mContext = RuntimeEnvironment.application;
         mShadowPackageManager =
-                (FrameworkShadowPackageManager)
-                        extract(RuntimeEnvironment.application.getPackageManager());
-        mContext = RuntimeEnvironment.application.getApplicationContext();
+                (FrameworkShadowPackageManager) extract(mContext.getPackageManager());
 
         mTransportA1 = genericTransport(PACKAGE_A, "TransportFoo");
         mTransportA2 = genericTransport(PACKAGE_A, "TransportBar");
         mTransportB1 = genericTransport(PACKAGE_B, "TransportBaz");
     }
 
-    @After
-    public void tearDown() throws Exception {
-        ShadowContextImplForBackup.resetBackupShadowState();
-    }
-
     @Test
     public void testRegisterTransports() throws Exception {
         setUpPackage(PACKAGE_A, ApplicationInfo.PRIVATE_FLAG_PRIVILEGED);
diff --git a/services/robotests/src/com/android/server/backup/restore/ActiveRestoreSessionTest.java b/services/robotests/src/com/android/server/backup/restore/ActiveRestoreSessionTest.java
new file mode 100644
index 0000000..4ac00f0
--- /dev/null
+++ b/services/robotests/src/com/android/server/backup/restore/ActiveRestoreSessionTest.java
@@ -0,0 +1,205 @@
+/*
+ * 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.restore;
+
+import static com.android.server.backup.testing.BackupManagerServiceTestUtils.createBackupWakeLock;
+import static com.android.server.backup.testing.BackupManagerServiceTestUtils.setUpBackupManagerServiceBasics;
+import static com.android.server.backup.testing.BackupManagerServiceTestUtils.startBackupThreadAndGetLooper;
+import static com.android.server.backup.testing.TransportData.backupTransport;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.AdditionalMatchers.aryEq;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+import static org.robolectric.Shadows.shadowOf;
+import static org.testng.Assert.expectThrows;
+
+import android.app.Application;
+import android.app.backup.IBackupManagerMonitor;
+import android.app.backup.IRestoreObserver;
+import android.app.backup.IRestoreSession;
+import android.app.backup.RestoreSet;
+import android.os.Looper;
+import android.os.PowerManager;
+import android.platform.test.annotations.Presubmit;
+
+import com.android.server.EventLogTags;
+import com.android.server.backup.BackupManagerService;
+import com.android.server.backup.TransportManager;
+import com.android.server.backup.internal.BackupHandler;
+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.testing.FrameworkRobolectricTestRunner;
+import com.android.server.testing.SystemLoaderPackages;
+import com.android.server.testing.shadows.ShadowEventLog;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+import org.robolectric.RuntimeEnvironment;
+import org.robolectric.annotation.Config;
+import org.robolectric.shadows.ShadowApplication;
+import org.robolectric.shadows.ShadowLooper;
+
+@RunWith(FrameworkRobolectricTestRunner.class)
+@Config(manifest = Config.NONE, sdk = 26, shadows = ShadowEventLog.class)
+@SystemLoaderPackages({"com.android.server.backup"})
+@Presubmit
+public class ActiveRestoreSessionTest {
+    private static final String PACKAGE_1 = "com.example.package1";
+    private static final String PACKAGE_2 = "com.example.package2";
+
+    @Mock private BackupManagerService mBackupManagerService;
+    @Mock private TransportManager mTransportManager;
+    @Mock private IRestoreObserver mObserver;
+    @Mock private IBackupManagerMonitor mMonitor;
+    private ShadowLooper mShadowBackupLooper;
+    private ShadowApplication mShadowApplication;
+    private PowerManager.WakeLock mWakeLock;
+    private TransportData mTransport;
+    private RestoreSet mRestoreSet1;
+    private RestoreSet mRestoreSet2;
+
+    @Before
+    public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
+        mTransport = backupTransport();
+
+        mRestoreSet1 = new RestoreSet("name1", "device1", 1L);
+        mRestoreSet2 = new RestoreSet("name2", "device2", 2L);
+
+        Application application = RuntimeEnvironment.application;
+        mShadowApplication = shadowOf(application);
+
+        Looper backupLooper = startBackupThreadAndGetLooper();
+        mShadowBackupLooper = shadowOf(backupLooper);
+        BackupHandler backupHandler = new BackupHandler(mBackupManagerService, backupLooper);
+
+        mWakeLock = createBackupWakeLock(application);
+
+        setUpBackupManagerServiceBasics(
+                mBackupManagerService,
+                application,
+                mTransportManager,
+                application.getPackageManager(),
+                backupHandler,
+                mWakeLock);
+    }
+
+    @Test
+    public void testGetAvailableRestoreSets_withoutPermission() throws Exception {
+        mShadowApplication.denyPermissions(android.Manifest.permission.BACKUP);
+        setUpTransport(mTransport);
+        IRestoreSession restoreSession = createActiveRestoreSession(PACKAGE_1, mTransport);
+
+        expectThrows(
+                SecurityException.class,
+                () -> restoreSession.getAvailableRestoreSets(mObserver, mMonitor));
+    }
+
+    @Test
+    public void testGetAvailableRestoreSets_forNullObserver() throws Exception {
+        mShadowApplication.grantPermissions(android.Manifest.permission.BACKUP);
+        setUpTransport(mTransport);
+        IRestoreSession restoreSession = createActiveRestoreSession(PACKAGE_1, mTransport);
+
+        expectThrows(
+                RuntimeException.class,
+                () -> restoreSession.getAvailableRestoreSets(null, mMonitor));
+    }
+
+    @Test
+    public void testGetAvailableRestoreSets_whenTransportNotRegistered() throws Exception {
+        mShadowApplication.grantPermissions(android.Manifest.permission.BACKUP);
+        setUpTransport(mTransport.unregistered());
+        IRestoreSession restoreSession = createActiveRestoreSession(PACKAGE_1, mTransport);
+
+        int result = restoreSession.getAvailableRestoreSets(mObserver, mMonitor);
+
+        assertThat(result).isEqualTo(-1);
+    }
+
+    @Test
+    public void testGetAvailableRestoreSets() throws Exception {
+        mShadowApplication.grantPermissions(android.Manifest.permission.BACKUP);
+        TransportMock transportMock = setUpTransport(mTransport);
+        when(transportMock.transport.getAvailableRestoreSets())
+                .thenReturn(new RestoreSet[] {mRestoreSet1, mRestoreSet2});
+        IRestoreSession restoreSession = createActiveRestoreSession(PACKAGE_1, mTransport);
+
+        int result = restoreSession.getAvailableRestoreSets(mObserver, mMonitor);
+
+        mShadowBackupLooper.runToEndOfTasks();
+        assertThat(result).isEqualTo(0);
+        verify(mObserver)
+                .restoreSetsAvailable(aryEq(new RestoreSet[] {mRestoreSet1, mRestoreSet2}));
+        verify(mTransportManager)
+                .disposeOfTransportClient(eq(transportMock.transportClient), any());
+        assertThat(mWakeLock.isHeld()).isFalse();
+    }
+
+    @Test
+    public void testGetAvailableRestoreSets_forEmptyRestoreSets() throws Exception {
+        mShadowApplication.grantPermissions(android.Manifest.permission.BACKUP);
+        TransportMock transportMock = setUpTransport(mTransport);
+        when(transportMock.transport.getAvailableRestoreSets()).thenReturn(new RestoreSet[0]);
+        IRestoreSession restoreSession = createActiveRestoreSession(PACKAGE_1, mTransport);
+
+        int result = restoreSession.getAvailableRestoreSets(mObserver, mMonitor);
+
+        mShadowBackupLooper.runToEndOfTasks();
+        assertThat(result).isEqualTo(0);
+        verify(mObserver).restoreSetsAvailable(aryEq(new RestoreSet[0]));
+        assertThat(ShadowEventLog.hasEvent(EventLogTags.RESTORE_TRANSPORT_FAILURE)).isFalse();
+    }
+
+    @Test
+    public void testGetAvailableRestoreSets_forNullRestoreSets() throws Exception {
+        mShadowApplication.grantPermissions(android.Manifest.permission.BACKUP);
+        TransportMock transportMock = setUpTransport(mTransport);
+        when(transportMock.transport.getAvailableRestoreSets()).thenReturn(null);
+        IRestoreSession restoreSession = createActiveRestoreSession(PACKAGE_1, mTransport);
+
+        int result = restoreSession.getAvailableRestoreSets(mObserver, mMonitor);
+
+        mShadowBackupLooper.runToEndOfTasks();
+        assertThat(result).isEqualTo(0);
+        verify(mObserver).restoreSetsAvailable(isNull());
+        assertThat(ShadowEventLog.hasEvent(EventLogTags.RESTORE_TRANSPORT_FAILURE)).isTrue();
+        verify(mTransportManager)
+                .disposeOfTransportClient(eq(transportMock.transportClient), any());
+        assertThat(mWakeLock.isHeld()).isFalse();
+    }
+
+    private IRestoreSession createActiveRestoreSession(
+            String packageName, TransportData transport) {
+        return new ActiveRestoreSession(
+                mBackupManagerService, packageName, transport.transportName);
+    }
+
+    private TransportMock setUpTransport(TransportData transport) throws Exception {
+        return TransportTestUtils.setUpTransport(mTransportManager, transport);
+    }
+}
diff --git a/services/robotests/src/com/android/server/backup/testing/BackupManagerServiceTestUtils.java b/services/robotests/src/com/android/server/backup/testing/BackupManagerServiceTestUtils.java
new file mode 100644
index 0000000..3c0234b
--- /dev/null
+++ b/services/robotests/src/com/android/server/backup/testing/BackupManagerServiceTestUtils.java
@@ -0,0 +1,95 @@
+/*
+ * 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.testing;
+
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.app.Application;
+import android.app.IActivityManager;
+import android.content.Context;
+import android.content.pm.PackageManager;
+import android.os.HandlerThread;
+import android.os.Looper;
+import android.os.PowerManager;
+import android.util.SparseArray;
+
+import com.android.server.backup.BackupManagerService;
+import com.android.server.backup.TransportManager;
+import com.android.server.backup.internal.BackupHandler;
+
+import java.lang.Thread.UncaughtExceptionHandler;
+
+public class BackupManagerServiceTestUtils {
+    /** Sets up a basic mocks for {@link BackupManagerService} */
+    public static void setUpBackupManagerServiceBasics(
+            BackupManagerService backupManagerService,
+            Context context,
+            TransportManager transportManager,
+            PackageManager packageManager,
+            BackupHandler backupHandler,
+            PowerManager.WakeLock wakeLock) {
+        when(backupManagerService.getContext()).thenReturn(context);
+        when(backupManagerService.getTransportManager()).thenReturn(transportManager);
+        when(backupManagerService.getPackageManager()).thenReturn(packageManager);
+        when(backupManagerService.getBackupHandler()).thenReturn(backupHandler);
+        when(backupManagerService.getCurrentOpLock()).thenReturn(new Object());
+        when(backupManagerService.getQueueLock()).thenReturn(new Object());
+        when(backupManagerService.getCurrentOperations()).thenReturn(new SparseArray<>());
+        when(backupManagerService.getActivityManager()).thenReturn(mock(IActivityManager.class));
+        when(backupManagerService.getWakelock()).thenReturn(wakeLock);
+    }
+
+    public static PowerManager.WakeLock createBackupWakeLock(Application application) {
+        PowerManager powerManager =
+                (PowerManager) application.getSystemService(Context.POWER_SERVICE);
+        return powerManager.newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, "*backup*");
+    }
+
+    /**
+     * Creates a backup thread associated with a looper, starts it and returns its looper for
+     * shadowing and creation of the backup handler.
+     *
+     * <p>Note that Robolectric simulates multi-thread in a single-thread to avoid flakiness, so
+     * even though we started the thread, you should control its execution via the shadow of the
+     * looper returned.
+     *
+     * @return The {@link Looper} for the backup thread.
+     */
+    public static Looper startBackupThreadAndGetLooper() {
+        HandlerThread backupThread = new HandlerThread("backup");
+        backupThread.start();
+        return backupThread.getLooper();
+    }
+
+    /**
+     * Similar to {@link #startBackupThreadAndGetLooper()} but with a custom exception handler and
+     * returning the thread instead of the looper associated with it.
+     *
+     * @param exceptionHandler Uncaught exception handler for backup thread.
+     * @return The backup thread.
+     * @see #startBackupThreadAndGetLooper()
+     */
+    public static HandlerThread startBackupThread(UncaughtExceptionHandler exceptionHandler) {
+        HandlerThread backupThread = new HandlerThread("backup");
+        backupThread.setUncaughtExceptionHandler(exceptionHandler);
+        backupThread.start();
+        return backupThread;
+    }
+
+    private BackupManagerServiceTestUtils() {}
+}
diff --git a/services/robotests/src/com/android/server/backup/testing/ShadowBackupTransportStub.java b/services/robotests/src/com/android/server/backup/testing/ShadowBackupTransportStub.java
deleted file mode 100644
index 48a12f0..0000000
--- a/services/robotests/src/com/android/server/backup/testing/ShadowBackupTransportStub.java
+++ /dev/null
@@ -1,40 +0,0 @@
-/*
- * Copyright (C) 2017 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.testing;
-
-import android.os.IBinder;
-
-import com.android.internal.backup.IBackupTransport;
-
-import org.robolectric.annotation.Implementation;
-import org.robolectric.annotation.Implements;
-
-import java.util.HashMap;
-import java.util.Map;
-
-/**
- * Shadow IBackupTransport.Stub, returns a transport corresponding to the binder.
- */
-@Implements(IBackupTransport.Stub.class)
-public class ShadowBackupTransportStub {
-    public static Map<IBinder, IBackupTransport> sBinderTransportMap = new HashMap<>();
-
-    @Implementation
-    public static IBackupTransport asInterface(IBinder obj) {
-        return sBinderTransportMap.get(obj);
-    }
-}
diff --git a/services/robotests/src/com/android/server/backup/testing/ShadowContextImplForBackup.java b/services/robotests/src/com/android/server/backup/testing/ShadowContextImplForBackup.java
deleted file mode 100644
index c3975db..0000000
--- a/services/robotests/src/com/android/server/backup/testing/ShadowContextImplForBackup.java
+++ /dev/null
@@ -1,64 +0,0 @@
-/*
- * Copyright (C) 2017 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.testing;
-
-import android.annotation.RequiresPermission;
-import android.content.ComponentName;
-import android.content.Intent;
-import android.content.ServiceConnection;
-import android.os.IBinder;
-import android.os.UserHandle;
-
-import org.robolectric.annotation.Implementation;
-import org.robolectric.annotation.Implements;
-import org.robolectric.shadows.ShadowApplication;
-import org.robolectric.shadows.ShadowContextImpl;
-
-import java.util.HashMap;
-import java.util.HashSet;
-import java.util.Map;
-import java.util.Set;
-
-/**
- * Implementation of ContextImpl shadow, handling bindServiceAsUser().
- */
-@Implements(className = ShadowContextImpl.CLASS_NAME)
-public class ShadowContextImplForBackup extends ShadowContextImpl {
-    public static Map<ComponentName, IBinder> sComponentBinderMap = new HashMap<>();
-    public static Set<ComponentName> sUnbindableComponents = new HashSet<>();
-
-    @Implementation
-    public boolean bindServiceAsUser(@RequiresPermission Intent service, ServiceConnection conn,
-            int flags, UserHandle user) {
-        if (sUnbindableComponents.contains(service.getComponent())) {
-            return false;
-        }
-
-        ShadowApplication.getInstance().setComponentNameAndServiceForBindService(
-                service.getComponent(), sComponentBinderMap.get(service.getComponent()));
-        return bindService(service, conn, flags);
-    }
-
-
-    /**
-     * Resets backup-related shadow state.
-     */
-    public static void resetBackupShadowState() {
-        sComponentBinderMap.clear();
-        sUnbindableComponents.clear();
-    }
-}
diff --git a/services/robotests/src/com/android/server/backup/testing/ShadowAppBackupUtils.java b/services/robotests/src/com/android/server/testing/shadows/ShadowAppBackupUtils.java
similarity index 96%
rename from services/robotests/src/com/android/server/backup/testing/ShadowAppBackupUtils.java
rename to services/robotests/src/com/android/server/testing/shadows/ShadowAppBackupUtils.java
index 73cb4c0..c3b2f0d 100644
--- a/services/robotests/src/com/android/server/backup/testing/ShadowAppBackupUtils.java
+++ b/services/robotests/src/com/android/server/testing/shadows/ShadowAppBackupUtils.java
@@ -14,7 +14,7 @@
  * limitations under the License
  */
 
-package com.android.server.backup.testing;
+package com.android.server.testing.shadows;
 
 import android.annotation.Nullable;
 import android.content.pm.PackageManager;
diff --git a/services/robotests/src/com/android/server/backup/testing/ShadowBackupPolicyEnforcer.java b/services/robotests/src/com/android/server/testing/shadows/ShadowBackupPolicyEnforcer.java
similarity index 62%
rename from services/robotests/src/com/android/server/backup/testing/ShadowBackupPolicyEnforcer.java
rename to services/robotests/src/com/android/server/testing/shadows/ShadowBackupPolicyEnforcer.java
index ef4764d..e76b9d1 100644
--- a/services/robotests/src/com/android/server/backup/testing/ShadowBackupPolicyEnforcer.java
+++ b/services/robotests/src/com/android/server/testing/shadows/ShadowBackupPolicyEnforcer.java
@@ -1,5 +1,6 @@
-package com.android.server.backup.testing;
+package com.android.server.testing.shadows;
 
+import android.annotation.Nullable;
 import android.content.ComponentName;
 
 import com.android.server.backup.BackupPolicyEnforcer;
@@ -9,14 +10,15 @@
 
 @Implements(BackupPolicyEnforcer.class)
 public class ShadowBackupPolicyEnforcer {
+    @Nullable private static ComponentName sMandatoryBackupTransport;
 
-    private static ComponentName sMandatoryBackupTransport;
-
-    public static void setMandatoryBackupTransport(ComponentName backupTransportComponent) {
+    public static void setMandatoryBackupTransport(
+            @Nullable ComponentName backupTransportComponent) {
         sMandatoryBackupTransport = backupTransportComponent;
     }
 
     @Implementation
+    @Nullable
     public ComponentName getMandatoryBackupTransport() {
         return sMandatoryBackupTransport;
     }