Merge "Fix the missing initialization of tracker"
diff --git a/device_build_interfaces/com/android/tradefed/device/TestDeviceOptions.java b/device_build_interfaces/com/android/tradefed/device/TestDeviceOptions.java
index 329de8e..0d10607 100644
--- a/device_build_interfaces/com/android/tradefed/device/TestDeviceOptions.java
+++ b/device_build_interfaces/com/android/tradefed/device/TestDeviceOptions.java
@@ -593,6 +593,11 @@
         return mGceDriverParams;
     }
 
+    /** Add a param to the gce driver params. */
+    public void addGceDriverParams(String param) {
+        mGceDriverParams.add(param);
+    }
+
     /** Set the GCE driver parameter that should be paired with the build id from build info */
     public void setGceDriverBuildIdParam(String gceDriverBuildIdParam) {
         mGceDriverBuildIdParam = gceDriverBuildIdParam;
diff --git a/invocation_interfaces/com/android/tradefed/invoker/logger/InvocationMetricLogger.java b/invocation_interfaces/com/android/tradefed/invoker/logger/InvocationMetricLogger.java
index a083a58..a23b5c1 100644
--- a/invocation_interfaces/com/android/tradefed/invoker/logger/InvocationMetricLogger.java
+++ b/invocation_interfaces/com/android/tradefed/invoker/logger/InvocationMetricLogger.java
@@ -37,7 +37,11 @@
         STAGE_TESTS_INDIVIDUAL_DOWNLOADS("stage_tests_individual_downloads", true),
         SHUTDOWN_HARD_LATENCY("shutdown_hard_latency_ms", false),
         DEVICE_DONE_TIMESTAMP("device_done_timestamp", false),
-        DEVICE_RELEASE_STATE("device_release_state", false);
+        DEVICE_RELEASE_STATE("device_release_state", false),
+        SANDBOX_EXIT_CODE("sandbox_exit_code", false),
+        CF_FETCH_ARTIFACT_TIME("cf_fetch_artifact_time_ms", false),
+        CF_GCE_CREATE_TIME("cf_gce_create_time_ms", false),
+        CF_LAUNCH_CVD_TIME("cf_launch_cvd_time_ms", false);
 
         private final String mKeyName;
         // Whether or not to add the value when the key is added again.
diff --git a/src/com/android/tradefed/build/FileDownloadCache.java b/src/com/android/tradefed/build/FileDownloadCache.java
index 1b8d68f..d79054e 100644
--- a/src/com/android/tradefed/build/FileDownloadCache.java
+++ b/src/com/android/tradefed/build/FileDownloadCache.java
@@ -328,6 +328,9 @@
             throws BuildRetrievalError {
         boolean download = false;
         File cachedFile, copyFile;
+        if (remotePath == null) {
+            throw new BuildRetrievalError("remote path was null.");
+        }
 
         lockFile(remotePath);
         try {
diff --git a/src/com/android/tradefed/device/cloud/GceAvdInfo.java b/src/com/android/tradefed/device/cloud/GceAvdInfo.java
index 5e9e5c4..81877f6 100644
--- a/src/com/android/tradefed/device/cloud/GceAvdInfo.java
+++ b/src/com/android/tradefed/device/cloud/GceAvdInfo.java
@@ -16,10 +16,13 @@
 package com.android.tradefed.device.cloud;
 
 import com.android.tradefed.command.remote.DeviceDescriptor;
+import com.android.tradefed.invoker.logger.InvocationMetricLogger;
+import com.android.tradefed.invoker.logger.InvocationMetricLogger.InvocationMetricKey;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.targetprep.TargetSetupError;
 import com.android.tradefed.util.FileUtil;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 import com.google.common.net.HostAndPort;
 
@@ -184,6 +187,7 @@
             if (devices != null) {
                 if (devices.length() == 1) {
                     JSONObject d = (JSONObject) devices.get(0);
+                    addCfStartTimeMetrics(d);
                     String ip = d.getString("ip");
                     String instanceName = d.getString("instance_name");
                     GceAvdInfo avdInfo =
@@ -225,4 +229,27 @@
         }
         return res;
     }
+
+    @VisibleForTesting
+    static void addCfStartTimeMetrics(JSONObject json) {
+        // These metrics may not be available for all GCE.
+        String fetch_artifact_time = json.optString("fetch_artifact_time");
+        if (!fetch_artifact_time.isEmpty()) {
+            InvocationMetricLogger.addInvocationMetrics(
+                    InvocationMetricKey.CF_FETCH_ARTIFACT_TIME,
+                    Double.valueOf(Double.parseDouble(fetch_artifact_time) * 1000).longValue());
+        }
+        String gce_create_time = json.optString("gce_create_time");
+        if (!gce_create_time.isEmpty()) {
+            InvocationMetricLogger.addInvocationMetrics(
+                    InvocationMetricKey.CF_GCE_CREATE_TIME,
+                    Double.valueOf(Double.parseDouble(gce_create_time) * 1000).longValue());
+        }
+        String launch_cvd_time = json.optString("launch_cvd_time");
+        if (!launch_cvd_time.isEmpty()) {
+            InvocationMetricLogger.addInvocationMetrics(
+                    InvocationMetricKey.CF_LAUNCH_CVD_TIME,
+                    Double.valueOf(Double.parseDouble(launch_cvd_time) * 1000).longValue());
+        }
+    }
 }
diff --git a/src/com/android/tradefed/device/recovery/BatteryRechargeDeviceRecovery.java b/src/com/android/tradefed/device/recovery/BatteryRechargeDeviceRecovery.java
index 85e4df0..fb268a0 100644
--- a/src/com/android/tradefed/device/recovery/BatteryRechargeDeviceRecovery.java
+++ b/src/com/android/tradefed/device/recovery/BatteryRechargeDeviceRecovery.java
@@ -36,7 +36,8 @@
             return true;
         }
         Integer level = device.getBattery();
-        if (level == null || level >= mMinBattery) {
+        // Skip zero battery since some devices use that as 'no-battery'
+        if (level == null || level == 0 || level >= mMinBattery) {
             return true;
         }
         return false;
diff --git a/src/com/android/tradefed/invoker/InvocationExecution.java b/src/com/android/tradefed/invoker/InvocationExecution.java
index 658808a..a1dcecf 100644
--- a/src/com/android/tradefed/invoker/InvocationExecution.java
+++ b/src/com/android/tradefed/invoker/InvocationExecution.java
@@ -331,6 +331,7 @@
     public final void runDevicePreInvocationSetup(
             IInvocationContext context, IConfiguration config, ITestLogger logger)
             throws DeviceNotAvailableException, TargetSetupError {
+        customizeDevicePreInvocation(config, context);
         for (String deviceName : context.getDeviceConfigNames()) {
             ITestDevice device = context.getDevice(deviceName);
 
@@ -347,6 +348,16 @@
         }
     }
 
+    /**
+     * Give a chance to customize some of the device before preInvocationSetup.
+     *
+     * @param config The config of the invocation.
+     * @param context The current invocation context.
+     */
+    protected void customizeDevicePreInvocation(IConfiguration config, IInvocationContext context) {
+        // Empty by default
+    }
+
     /** {@inheritDoc} */
     @Override
     public final void runDevicePostInvocationTearDown(
diff --git a/src/com/android/tradefed/invoker/RemoteInvocationExecution.java b/src/com/android/tradefed/invoker/RemoteInvocationExecution.java
index 99d9ab2..56fdb7c 100644
--- a/src/com/android/tradefed/invoker/RemoteInvocationExecution.java
+++ b/src/com/android/tradefed/invoker/RemoteInvocationExecution.java
@@ -30,6 +30,7 @@
 import com.android.tradefed.config.OptionSetter;
 import com.android.tradefed.device.DeviceNotAvailableException;
 import com.android.tradefed.device.DeviceSelectionOptions;
+import com.android.tradefed.device.ITestDevice;
 import com.android.tradefed.device.TestDeviceOptions;
 import com.android.tradefed.device.cloud.GceAvdInfo;
 import com.android.tradefed.device.cloud.GceManager;
@@ -117,6 +118,21 @@
     }
 
     @Override
+    protected void customizeDevicePreInvocation(IConfiguration config, IInvocationContext context) {
+        super.customizeDevicePreInvocation(config, context);
+
+        if (config.getCommandOptions().getShardCount() != null
+                && config.getCommandOptions().getShardIndex() == null) {
+            ITestDevice device = context.getDevices().get(0);
+            TestDeviceOptions options = device.getOptions();
+            // Trigger the multi-tenant start in the VM
+            options.addGceDriverParams("--num-avds-per-instance");
+            options.addGceDriverParams(config.getCommandOptions().getShardCount().toString());
+            // TODO: Track how many instances we created
+        }
+    }
+
+    @Override
     public void runTests(
             TestInformation info, IConfiguration config, ITestInvocationListener listener)
             throws Throwable {
diff --git a/src/com/android/tradefed/sandbox/SandboxInvocationRunner.java b/src/com/android/tradefed/sandbox/SandboxInvocationRunner.java
index 96cb4f7..e90ff43 100644
--- a/src/com/android/tradefed/sandbox/SandboxInvocationRunner.java
+++ b/src/com/android/tradefed/sandbox/SandboxInvocationRunner.java
@@ -18,6 +18,8 @@
 import com.android.tradefed.config.Configuration;
 import com.android.tradefed.config.IConfiguration;
 import com.android.tradefed.invoker.TestInformation;
+import com.android.tradefed.invoker.logger.InvocationMetricLogger;
+import com.android.tradefed.invoker.logger.InvocationMetricLogger.InvocationMetricKey;
 import com.android.tradefed.log.LogUtil.CLog;
 import com.android.tradefed.result.ITestInvocationListener;
 import com.android.tradefed.util.CommandResult;
@@ -60,6 +62,11 @@
         PrettyPrintDelimiter.printStageDelimiter("Done with Sandbox Environment Setup");
         try {
             CommandResult result = sandbox.run(config, listener);
+            if (result.getExitCode() != null) {
+                // Log the exit code
+                InvocationMetricLogger.addInvocationMetrics(
+                        InvocationMetricKey.SANDBOX_EXIT_CODE, result.getExitCode());
+            }
             if (!CommandStatus.SUCCESS.equals(result.getStatus())) {
                 CLog.e(
                         "Sandbox finished with status: %s and exit code: %s",
diff --git a/src/com/android/tradefed/testtype/DeviceBatteryLevelChecker.java b/src/com/android/tradefed/testtype/DeviceBatteryLevelChecker.java
index f097786..14f8e78 100644
--- a/src/com/android/tradefed/testtype/DeviceBatteryLevelChecker.java
+++ b/src/com/android/tradefed/testtype/DeviceBatteryLevelChecker.java
@@ -92,6 +92,11 @@
         mTestDevice.executeShellCommand("stop");
     }
 
+    private void startDeviceRuntime() throws DeviceNotAvailableException {
+        mTestDevice.executeShellCommand("start");
+        mTestDevice.waitForDeviceAvailable();
+    }
+
     /** {@inheritDoc} */
     @Override
     public void run(TestInformation testInfo, ITestInvocationListener listener)
@@ -171,6 +176,12 @@
             }
             batteryLevel = newLevel;
         }
+
+        if (mStopRuntime) {
+            // Restart runtime if it was stopped
+            startDeviceRuntime();
+        }
+
         CLog.w("Device %s is now charged to battery level %d; releasing.",
                 mTestDevice.getSerialNumber(), batteryLevel);
     }
diff --git a/test_framework/com/android/tradefed/targetprep/InstallApexModuleTargetPreparer.java b/test_framework/com/android/tradefed/targetprep/InstallApexModuleTargetPreparer.java
index 199037b..8e3ef8c 100644
--- a/test_framework/com/android/tradefed/targetprep/InstallApexModuleTargetPreparer.java
+++ b/test_framework/com/android/tradefed/targetprep/InstallApexModuleTargetPreparer.java
@@ -87,6 +87,13 @@
 
         cleanUpStagedAndActiveSession(device);
 
+        Set<ApexInfo> activatedApexes = device.getActiveApexes();
+
+        CLog.i("Activated apex packages list before module/train installation:");
+        for (ApexInfo info : activatedApexes) {
+            CLog.i("Activated apex: %s", info.toString());
+        }
+
         List<String> testAppFileNames = getModulesToInstall(testInfo);
         if (testAppFileNames.isEmpty()) {
             CLog.i("No modules are preloaded on the device, so no modules will be installed.");
@@ -114,7 +121,7 @@
             }
         }
 
-        Set<ApexInfo> activatedApexes = device.getActiveApexes();
+        activatedApexes = device.getActiveApexes();
 
         if (activatedApexes.isEmpty()) {
             throw new TargetSetupError(
@@ -122,6 +129,11 @@
                             "Failed to retrieve activated apex on device %s. Empty set returned.",
                             device.getSerialNumber()),
                     device.getDeviceDescriptor());
+        } else {
+            CLog.i("Activated apex packages list after module/train installation:");
+            for (ApexInfo info : activatedApexes) {
+                CLog.i("Activated apex: %s", info.toString());
+            }
         }
 
         List<ApexInfo> failToActivateApex = new ArrayList<ApexInfo>();
@@ -133,10 +145,6 @@
         }
 
         if (!failToActivateApex.isEmpty()) {
-            CLog.i("Activated apex packages list:");
-            for (ApexInfo info : activatedApexes) {
-                CLog.i("Activated apex: %s", info.toString());
-            }
             throw new TargetSetupError(
                     String.format(
                             "Failed to activate %s on device %s.",
diff --git a/tests/src/com/android/tradefed/build/FileDownloadCacheTest.java b/tests/src/com/android/tradefed/build/FileDownloadCacheTest.java
index d0643c5..6e84b35 100644
--- a/tests/src/com/android/tradefed/build/FileDownloadCacheTest.java
+++ b/tests/src/com/android/tradefed/build/FileDownloadCacheTest.java
@@ -92,6 +92,18 @@
         EasyMock.verify(mMockDownloader);
     }
 
+    @Test
+    public void testFetchRemoteFile_destFile_nullPath() throws Exception {
+        EasyMock.replay(mMockDownloader);
+        try {
+            assertFetchRemoteFile(null, null, null);
+            fail("Should have thrown an exception.");
+        } catch (BuildRetrievalError expected) {
+            assertEquals("remote path was null.", expected.getMessage());
+        }
+        EasyMock.verify(mMockDownloader);
+    }
+
     /**
      * Test {@link FileDownloadCache#fetchRemoteFile(IFileDownloader, String)} when file can be
      * retrieved from cache.
diff --git a/tests/src/com/android/tradefed/config/DynamicRemoteFileResolverTest.java b/tests/src/com/android/tradefed/config/DynamicRemoteFileResolverTest.java
index c9d8c38..3cf9e8a 100644
--- a/tests/src/com/android/tradefed/config/DynamicRemoteFileResolverTest.java
+++ b/tests/src/com/android/tradefed/config/DynamicRemoteFileResolverTest.java
@@ -446,7 +446,13 @@
         List<Set<File>> downloadedFile = null;
         try {
             downloadedFile = executor.invokeAll(call, 1, TimeUnit.MINUTES);
-            assertEquals(3, downloadedFile.get(0).size());
+            boolean oneMustBeNonEmpty = false;
+            for (Set<File> set : downloadedFile) {
+                if (set.size() == 3) {
+                    oneMustBeNonEmpty = true;
+                }
+            }
+            assertTrue(oneMustBeNonEmpty);
             // The file has been replaced by the downloaded one.
             assertEquals(1, object.remoteMultiMap.size());
             assertEquals(3, object.remoteMultiMap.values().size());
diff --git a/tests/src/com/android/tradefed/device/cloud/GceAvdInfoTest.java b/tests/src/com/android/tradefed/device/cloud/GceAvdInfoTest.java
index 79e7c15..35ed79d 100644
--- a/tests/src/com/android/tradefed/device/cloud/GceAvdInfoTest.java
+++ b/tests/src/com/android/tradefed/device/cloud/GceAvdInfoTest.java
@@ -21,11 +21,17 @@
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 
+import com.android.tradefed.invoker.logger.InvocationMetricLogger;
+import com.android.tradefed.invoker.logger.InvocationMetricLogger.InvocationMetricKey;
 import com.android.tradefed.targetprep.TargetSetupError;
 
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
+import org.json.JSONObject;
+import org.json.JSONArray;
+
+import java.util.Map;
 
 /** Unit tests for {@link GceAvdInfo} */
 @RunWith(JUnit4.class)
@@ -306,4 +312,33 @@
             // expected
         }
     }
+
+    /** Test CF start time metrics are added. */
+    @Test
+    public void testCfStartTimeMetricsAdded() throws Exception {
+        String cuttlefish =
+                " {\n"
+                        + "    \"command\": \"create_cf\",\n"
+                        + "    \"data\": {\n"
+                        + "      \"devices\": [\n"
+                        + "        {\n"
+                        + "          \"ip\": \"34.71.83.182\",\n"
+                        + "          \"instance_name\": \"ins-cf-x86-phone-userdebug\",\n"
+                        + "          \"fetch_artifact_time\": 63.22,\n"
+                        + "          \"gce_create_time\": 23.5,\n"
+                        + "          \"launch_cvd_time\": 226.5\n"
+                        + "        },\n"
+                        + "      ]\n"
+                        + "    },\n"
+                        + "    \"errors\": [],\n"
+                        + "    \"status\": \"SUCCESS\"\n"
+                        + "  }";
+        JSONObject res = new JSONObject(cuttlefish);
+        JSONArray devices = res.getJSONObject("data").getJSONArray("devices");
+        GceAvdInfo.addCfStartTimeMetrics((JSONObject) devices.get(0));
+        Map<String, String> metrics = InvocationMetricLogger.getInvocationMetrics();
+        assertEquals("63220", metrics.get(InvocationMetricKey.CF_FETCH_ARTIFACT_TIME.toString()));
+        assertEquals("23500", metrics.get(InvocationMetricKey.CF_GCE_CREATE_TIME.toString()));
+        assertEquals("226500", metrics.get(InvocationMetricKey.CF_LAUNCH_CVD_TIME.toString()));
+    }
 }
diff --git a/tests/src/com/android/tradefed/invoker/SandboxedInvocationExecutionTest.java b/tests/src/com/android/tradefed/invoker/SandboxedInvocationExecutionTest.java
index c6beebd..654055e 100644
--- a/tests/src/com/android/tradefed/invoker/SandboxedInvocationExecutionTest.java
+++ b/tests/src/com/android/tradefed/invoker/SandboxedInvocationExecutionTest.java
@@ -251,6 +251,7 @@
                 .saveLogData(any(), any(), any());
 
         CommandResult result = new CommandResult(CommandStatus.SUCCESS);
+        result.setExitCode(0);
         doReturn(result).when(mMockSandbox).run(any(), any());
 
         doReturn(new BuildInfo()).when(mMockProvider).getBuild();
diff --git a/tests/src/com/android/tradefed/targetprep/InstallApexModuleTargetPreparerTest.java b/tests/src/com/android/tradefed/targetprep/InstallApexModuleTargetPreparerTest.java
index bd1f785..83cd35b 100644
--- a/tests/src/com/android/tradefed/targetprep/InstallApexModuleTargetPreparerTest.java
+++ b/tests/src/com/android/tradefed/targetprep/InstallApexModuleTargetPreparerTest.java
@@ -232,7 +232,7 @@
         mockSuccessfulInstallPackageAndReboot(mFakeApex);
         Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
         activatedApex.add(new ApexInfo("com.android.FAKE_APEX_PACKAGE_NAME", 1));
-        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(3);
         Set<String> installableModules = new HashSet<>();
         installableModules.add(APEX_PACKAGE_NAME);
         EasyMock.expect(mMockDevice.getInstalledPackageNames()).andReturn(installableModules);
@@ -253,7 +253,7 @@
         mockSuccessfulInstallPackageAndReboot(mFakeApex);
         Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
         activatedApex.add(new ApexInfo("com.android.FAKE_APEX_PACKAGE_NAME", 1));
-        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(3);
         Set<String> installableModules = new HashSet<>();
         installableModules.add(APEX_PACKAGE_NAME);
         EasyMock.expect(mMockDevice.getInstalledPackageNames()).andReturn(installableModules);
@@ -281,7 +281,7 @@
         mockSuccessfulInstallPackageAndReboot(mFakeApex);
         Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
         activatedApex.add(new ApexInfo("com.android.FAKE_APEX_PACKAGE_NAME", 1));
-        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(3);
         Set<String> installableModules = new HashSet<>();
         installableModules.add(APEX_PACKAGE_NAME);
         EasyMock.expect(mMockDevice.getInstalledPackageNames()).andReturn(installableModules);
@@ -308,7 +308,7 @@
         mMockDevice.reboot();
         EasyMock.expectLastCall();
         mockSuccessfulInstallPackageAndReboot(mFakeApex);
-        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(new HashSet<ApexInfo>()).times(2);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(new HashSet<ApexInfo>()).times(3);
         Set<String> installableModules = new HashSet<>();
         installableModules.add(APEX_PACKAGE_NAME);
         EasyMock.expect(mMockDevice.getInstalledPackageNames()).andReturn(installableModules);
@@ -345,7 +345,7 @@
         mockSuccessfulInstallPackageAndReboot(mFakeApex);
         Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
         activatedApex.add(new ApexInfo("com.android.FAKE_APEX_PACKAGE_NAME_TO_FAIL", 1));
-        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(3);
         Set<String> installableModules = new HashSet<>();
         installableModules.add(APEX_PACKAGE_NAME);
         EasyMock.expect(mMockDevice.getInstalledPackageNames()).andReturn(installableModules);
@@ -391,6 +391,7 @@
                 .andReturn(null)
                 .once();
         EasyMock.expect(mMockDevice.uninstallPackage(APK_PACKAGE_NAME)).andReturn(null).once();
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(new HashSet<ApexInfo>()).times(1);
         Set<String> installableModules = new HashSet<>();
         installableModules.add(APK_PACKAGE_NAME);
         EasyMock.expect(mMockDevice.getInstalledPackageNames()).andReturn(installableModules);
@@ -425,6 +426,7 @@
         mockSuccessfulInstallMultiApkWithoutReboot(apks);
         EasyMock.expect(mMockDevice.uninstallPackage(APK_PACKAGE_NAME)).andReturn(null).once();
         EasyMock.expect(mMockDevice.uninstallPackage(APK2_PACKAGE_NAME)).andReturn(null).once();
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(new HashSet<ApexInfo>()).times(1);
         Set<String> installableModules = new HashSet<>();
         installableModules.add(APK_PACKAGE_NAME);
         installableModules.add(APK2_PACKAGE_NAME);
@@ -470,6 +472,7 @@
         EasyMock.expect(mMockDevice.uninstallPackage(PERSISTENT_APK_PACKAGE_NAME))
                 .andReturn(null)
                 .once();
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(new HashSet<ApexInfo>()).times(1);
         Set<String> installableModules = new HashSet<>();
         installableModules.add(APK_PACKAGE_NAME);
         installableModules.add(APK2_PACKAGE_NAME);
@@ -547,6 +550,9 @@
             EasyMock.expect(mMockDevice.uninstallPackage(SPLIT_APK_PACKAGE_NAME))
                 .andReturn(null)
                 .once();
+            EasyMock.expect(mMockDevice.getActiveApexes())
+                    .andReturn(new HashSet<ApexInfo>())
+                    .times(1);
             Set<String> installableModules = new HashSet<>();
             installableModules.add(APK_PACKAGE_NAME);
             installableModules.add(SPLIT_APK_PACKAGE_NAME);
@@ -596,7 +602,7 @@
         mockSuccessfulInstallPackageAndReboot(mFakeApex);
         Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
         activatedApex.add(new ApexInfo("com.android.FAKE_APEX_PACKAGE_NAME", 1));
-        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(3);
         mMockDevice.reboot();
         EasyMock.expectLastCall();
         Set<String> installableModules = new HashSet<>();
@@ -645,7 +651,7 @@
         mockSuccessfulInstallMultiPackageAndReboot();
         Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
         activatedApex.add(new ApexInfo("com.android.FAKE_APEX_PACKAGE_NAME", 1));
-        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(3);
         EasyMock.expect(mMockDevice.uninstallPackage(APK_PACKAGE_NAME)).andReturn(null).once();
         mMockDevice.reboot();
         EasyMock.expectLastCall();
@@ -742,7 +748,7 @@
             mMockDevice.reboot();
             Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
             activatedApex.add(new ApexInfo(SPLIT_APEX_PACKAGE_NAME, 1));
-            EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
+            EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(3);
             EasyMock.expect(mMockDevice.uninstallPackage(SPLIT_APK_PACKAGE_NAME))
                 .andReturn(null)
                 .once();
@@ -845,7 +851,7 @@
 
             Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
             activatedApex.add(new ApexInfo(SPLIT_APEX_PACKAGE_NAME, 1));
-            EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(1);
+            EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
             EasyMock.expect(mMockDevice.uninstallPackage(SPLIT_APK_PACKAGE_NAME))
                     .andReturn(null)
                     .once();
@@ -947,6 +953,7 @@
         EasyMock.expect(mMockDevice.executeShellV2Command("ls " + APEX_DATA_DIR)).andReturn(res);
         EasyMock.expect(mMockDevice.executeShellV2Command("ls " + SESSION_DATA_DIR)).andReturn(res);
         EasyMock.expect(mMockDevice.executeShellV2Command("ls " + STAGING_DATA_DIR)).andReturn(res);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(new HashSet<ApexInfo>()).times(1);
         mMockDevice.reboot();
         EasyMock.expectLastCall();
 
@@ -977,7 +984,7 @@
         mockSuccessfulInstallMultiPackageAndReboot();
         Set<ApexInfo> activatedApex = new HashSet<ApexInfo>();
         activatedApex.add(new ApexInfo("com.android.FAKE_APEX_PACKAGE_NAME", 1));
-        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(2);
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(3);
         EasyMock.expect(mMockDevice.uninstallPackage(APK_PACKAGE_NAME)).andReturn(null).once();
         mMockDevice.reboot();
         EasyMock.expectLastCall();
@@ -1037,6 +1044,7 @@
                 .andReturn(null)
                 .once();
         EasyMock.expect(mMockDevice.uninstallPackage(APK_PACKAGE_NAME)).andReturn(null).once();
+        EasyMock.expect(mMockDevice.getActiveApexes()).andReturn(activatedApex).times(1);
 
         EasyMock.replay(mMockBuildInfo, mMockDevice);
         mInstallApexModuleTargetPreparer.setUp(mTestInfo);