Ensure necessary reboot to happen in tests

In one case, if failure happens during the preparation, fail the test
after the reboot.

Bug: 279127811
Test: atest odsign_e2e_tests
Change-Id: I42c68254dd4d0060de01e3e5fd7b0b5a4f13df3b
diff --git a/test/odsign/test-src/com/android/tests/odsign/CompOsDenialHostTest.java b/test/odsign/test-src/com/android/tests/odsign/CompOsDenialHostTest.java
index 680a976..535cd91 100644
--- a/test/odsign/test-src/com/android/tests/odsign/CompOsDenialHostTest.java
+++ b/test/odsign/test-src/com/android/tests/odsign/CompOsDenialHostTest.java
@@ -74,14 +74,16 @@
     public static void afterClassWithDevice(TestInformation testInfo) throws Exception {
         OdsignTestUtils testUtils = new OdsignTestUtils(testInfo);
 
-        // Remove all test states.
-        testInfo.getDevice().executeShellV2Command("rm -rf " + PENDING_ARTIFACTS_DIR);
-        testInfo.getDevice().executeShellV2Command("rm -rf " + PENDING_ARTIFACTS_BACKUP_DIR);
-        testUtils.removeCompilationLogToAvoidBackoff();
-        testUtils.uninstallTestApex();
-
-        // Reboot should restore the device back to a good state.
-        testUtils.reboot();
+        try {
+            // Remove all test states.
+            testInfo.getDevice().executeShellV2Command("rm -rf " + PENDING_ARTIFACTS_DIR);
+            testInfo.getDevice().executeShellV2Command("rm -rf " + PENDING_ARTIFACTS_BACKUP_DIR);
+            testUtils.removeCompilationLogToAvoidBackoff();
+            testUtils.uninstallTestApex();
+        } finally {
+            // Reboot should restore the device back to a good state.
+            testUtils.reboot();
+        }
     }
 
     @Before
diff --git a/test/odsign/test-src/com/android/tests/odsign/CompOsSigningHostTest.java b/test/odsign/test-src/com/android/tests/odsign/CompOsSigningHostTest.java
index 5660ce9..cc56736 100644
--- a/test/odsign/test-src/com/android/tests/odsign/CompOsSigningHostTest.java
+++ b/test/odsign/test-src/com/android/tests/odsign/CompOsSigningHostTest.java
@@ -61,26 +61,32 @@
                 compOsTestUtils.checksumDirectoryContentPartial(
                         OdsignTestUtils.ART_APEX_DALVIK_CACHE_DIRNAME));
 
-        testUtils.installTestApex();
+        try {
+            testUtils.installTestApex();
 
-        testInfo.properties().put(TIMESTAMP_VM_START_KEY,
-                        String.valueOf(testUtils.getCurrentTimeMs()));
+            testInfo.properties().put(TIMESTAMP_VM_START_KEY,
+                    String.valueOf(testUtils.getCurrentTimeMs()));
 
-        compOsTestUtils.runCompilationJobEarlyAndWait();
+            compOsTestUtils.runCompilationJobEarlyAndWait();
 
-        testInfo.properties().put(PENDING_CHECKSUMS_KEY,
-                compOsTestUtils.checksumDirectoryContentPartial(PENDING_ARTIFACTS_DIR));
-
-        testInfo.properties().put(TIMESTAMP_REBOOT_KEY,
-                        String.valueOf(testUtils.getCurrentTimeMs()));
-        testUtils.reboot();
+            testInfo.properties().put(PENDING_CHECKSUMS_KEY,
+                    compOsTestUtils.checksumDirectoryContentPartial(PENDING_ARTIFACTS_DIR));
+            testInfo.properties().put(TIMESTAMP_REBOOT_KEY,
+                    String.valueOf(testUtils.getCurrentTimeMs()));
+        } finally {
+            // Make sure to reboot after installTestApex.
+            testUtils.reboot();
+        }
     }
 
     @AfterClassWithInfo
     public static void afterClassWithDevice(TestInformation testInfo) throws Exception {
         OdsignTestUtils testUtils = new OdsignTestUtils(testInfo);
-        testUtils.uninstallTestApex();
-        testUtils.reboot();
+        try {
+            testUtils.uninstallTestApex();
+        } finally {
+            testUtils.reboot();  // must happen
+        }
     }
 
     @Test