diff --git a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
index f9ece56..2ae0d6b 100644
--- a/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
+++ b/tests/testapk/src/java/com/android/microdroid/test/MicrodroidTests.java
@@ -93,7 +93,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicReference;
 
-import androidx.test.core.app.ApplicationProvider;
 import co.nstant.in.cbor.CborDecoder;
 import co.nstant.in.cbor.model.Array;
 import co.nstant.in.cbor.model.DataItem;
@@ -157,7 +156,7 @@
                             tr.mApkContentsPath = ts.getApkContentsPath();
                             tr.mEncryptedStoragePath = ts.getEncryptedStoragePath();
                         });
-        assertThat(testResults.mException).isNull();
+        testResults.assertNoException();
         assertThat(testResults.mAddInteger).isEqualTo(123 + 456);
         assertThat(testResults.mAppRunProp).isEqualTo("true");
         assertThat(testResults.mSublibRunProp).isEqualTo("true");
@@ -182,12 +181,8 @@
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm", config);
 
         TestResults testResults =
-                runVmTestService(
-                        vm,
-                        (ts, tr) -> {
-                            tr.mAddInteger = ts.addInteger(37, 73);
-                        });
-        assertThat(testResults.mException).isNull();
+                runVmTestService(vm, (ts, tr) -> tr.mAddInteger = ts.addInteger(37, 73));
+        testResults.assertNoException();
         assertThat(testResults.mAddInteger).isEqualTo(37 + 73);
     }
 
@@ -364,19 +359,14 @@
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_vsock", config);
 
-        AtomicReference<Exception> exception = new AtomicReference<>();
         AtomicReference<String> response = new AtomicReference<>();
         String request = "Look not into the abyss";
 
-        VmEventListener listener =
-                new VmEventListener() {
-                    @Override
-                    public void onPayloadReady(VirtualMachine vm) {
-                        try (vm) {
-                            ITestService testService =
-                                    ITestService.Stub.asInterface(
-                                            vm.connectToVsockServer(ITestService.SERVICE_PORT));
-                            testService.runEchoReverseServer();
+        TestResults testResults =
+                runVmTestService(
+                        vm,
+                        (service, results) -> {
+                            service.runEchoReverseServer();
 
                             ParcelFileDescriptor pfd =
                                     vm.connectVsock(ITestService.ECHO_REVERSE_PORT);
@@ -389,15 +379,8 @@
                                 writer.flush();
                                 response.set(reader.readLine());
                             }
-                        } catch (Exception e) {
-                            exception.set(e);
-                        }
-                    }
-                };
-        listener.runToFinish(TAG, vm);
-        if (exception.get() != null) {
-            throw new RuntimeException(exception.get());
-        }
+                        });
+        testResults.assertNoException();
         assertThat(response.get()).isEqualTo(new StringBuilder(request).reverse().toString());
     }
 
@@ -421,7 +404,6 @@
 
         // Maximal has everything that can be set to some non-default value. (And has different
         // values than minimal for the required fields.)
-        int maxCpus = Runtime.getRuntime().availableProcessors();
         VirtualMachineConfig.Builder maximalBuilder =
                 new VirtualMachineConfig.Builder(getContext())
                         .setProtectedVm(mProtectedVm)
@@ -783,7 +765,7 @@
                         (ts, tr) -> {
                             tr.mApkContentsPath = ts.getApkContentsPath();
                         });
-        assertThat(testResults.mException).isNull();
+        testResults.assertNoException();
         assertThat(testResults.mApkContentsPath).isEqualTo("/mnt/apk");
     }
 
@@ -906,15 +888,16 @@
 
     private VmCdis launchVmAndGetCdis(String instanceName) throws Exception {
         VirtualMachine vm = getVirtualMachineManager().get(instanceName);
-        final VmCdis vmCdis = new VmCdis();
-        final CompletableFuture<Exception> exception = new CompletableFuture<>();
+        VmCdis vmCdis = new VmCdis();
+        CompletableFuture<Exception> exception = new CompletableFuture<>();
         VmEventListener listener =
                 new VmEventListener() {
                     @Override
                     public void onPayloadReady(VirtualMachine vm) {
                         try {
-                            ITestService testService = ITestService.Stub.asInterface(
-                                    vm.connectToVsockServer(ITestService.SERVICE_PORT));
+                            ITestService testService =
+                                    ITestService.Stub.asInterface(
+                                            vm.connectToVsockServer(ITestService.SERVICE_PORT));
                             vmCdis.cdiAttest = testService.insecurelyExposeAttestationCdi();
                             vmCdis.instanceSecret = testService.insecurelyExposeVmInstanceSecret();
                         } catch (Exception e) {
@@ -998,26 +981,14 @@
                         .setDebugLevel(DEBUG_LEVEL_FULL)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("bcc_vm", normalConfig);
-        final CompletableFuture<byte[]> bcc = new CompletableFuture<>();
-        final CompletableFuture<Exception> exception = new CompletableFuture<>();
-        VmEventListener listener =
-                new VmEventListener() {
-                    @Override
-                    public void onPayloadReady(VirtualMachine vm) {
-                        try {
-                            ITestService testService = ITestService.Stub.asInterface(
-                                    vm.connectToVsockServer(ITestService.SERVICE_PORT));
-                            bcc.complete(testService.getBcc());
-                        } catch (Exception e) {
-                            exception.complete(e);
-                        } finally {
-                            forceStop(vm);
-                        }
-                    }
-                };
-        listener.runToFinish(TAG, vm);
-        byte[] bccBytes = bcc.getNow(null);
-        assertThat(exception.getNow(null)).isNull();
+        TestResults testResults =
+                runVmTestService(
+                        vm,
+                        (service, results) -> {
+                            results.mBcc = service.getBcc();
+                        });
+        testResults.assertNoException();
+        byte[] bccBytes = testResults.mBcc;
         assertThat(bccBytes).isNotNull();
 
         ByteArrayInputStream bais = new ByteArrayInputStream(bccBytes);
@@ -1054,10 +1025,6 @@
 
     private static final UUID MICRODROID_PARTITION_UUID =
             UUID.fromString("cf9afe9a-0662-11ec-a329-c32663a09d75");
-    private static final UUID U_BOOT_AVB_PARTITION_UUID =
-            UUID.fromString("7e8221e7-03e6-4969-948b-73a4c809a4f2");
-    private static final UUID U_BOOT_ENV_PARTITION_UUID =
-            UUID.fromString("0ab72d30-86ae-4d05-81b2-c1760be2b1f9");
     private static final UUID PVM_FW_PARTITION_UUID =
             UUID.fromString("90d2174a-038a-4bc6-adf3-824848fc5825");
     private static final long BLOCK_SIZE = 512;
@@ -1306,7 +1273,7 @@
                             tr.mAddInteger = ts.addInteger(123, 456);
                             tr.mEncryptedStoragePath = ts.getEncryptedStoragePath();
                         });
-        assertThat(origTestResults.mException).isNull();
+        origTestResults.assertNoException();
         assertThat(origTestResults.mAddInteger).isEqualTo(123 + 456);
         VirtualMachineManager vmm = getVirtualMachineManager();
         if (vmm.get(vmNameImport) != null) {
@@ -1332,7 +1299,7 @@
                             tr.mAddInteger = ts.addInteger(123, 456);
                             tr.mEncryptedStoragePath = ts.getEncryptedStoragePath();
                         });
-        assertThat(testResults.mException).isNull();
+        testResults.assertNoException();
         assertThat(testResults.mAddInteger).isEqualTo(123 + 456);
         return testResults;
     }
@@ -1384,7 +1351,7 @@
                                     /* content= */ EXAMPLE_STRING,
                                     /* path= */ "/mnt/encryptedstore/test_file");
                         });
-        assertThat(testResults.mException).isNull();
+        testResults.assertNoException();
 
         // Start a different vm (this changes the vm identity)
         VirtualMachine diff_test_vm = forceCreateNewVirtualMachine("diff_test_vm", config);
@@ -1439,7 +1406,7 @@
                             tr.mEffectiveCapabilities = ts.getEffectiveCapabilities();
                         });
 
-        assertThat(testResults.mException).isNull();
+        testResults.assertNoException();
         assertThat(testResults.mEffectiveCapabilities).isEmpty();
     }
 
@@ -1464,7 +1431,7 @@
                                     /* content= */ EXAMPLE_STRING,
                                     /* path= */ "/mnt/encryptedstore/test_file");
                         });
-        assertThat(testResults.mException).isNull();
+        testResults.assertNoException();
 
         // Re-run the same VM & verify the file persisted. Note, the previous `runVmTestService`
         // stopped the VM
@@ -1474,7 +1441,7 @@
                         (ts, tr) -> {
                             tr.mFileContent = ts.readFromFile("/mnt/encryptedstore/test_file");
                         });
-        assertThat(testResults.mException).isNull();
+        testResults.assertNoException();
         assertThat(testResults.mFileContent).isEqualTo(EXAMPLE_STRING);
     }
 
@@ -1498,7 +1465,7 @@
                             ts.mFileContent = testService.readFromFile("/mnt/apk/assets/file.txt");
                         });
 
-        assertThat(testResults.mException).isNull();
+        testResults.assertNoException();
         assertThat(testResults.mFileContent).isEqualTo("Hello, I am a file!");
     }
 
@@ -1507,7 +1474,7 @@
         assumeSupportedKernel();
 
         final VirtualMachineConfig vmConfig =
-                new VirtualMachineConfig.Builder(ApplicationProvider.getApplicationContext())
+                new VirtualMachineConfig.Builder(getContext())
                         .setProtectedVm(mProtectedVm)
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
                         .setDebugLevel(DEBUG_LEVEL_FULL)
@@ -1529,7 +1496,7 @@
         String time =
                 LocalDateTime.now().format(DateTimeFormatter.ofPattern("yyyy-MM-dd HH:mm:ss.SSS"));
         final VirtualMachineConfig vmConfig =
-                new VirtualMachineConfig.Builder(ApplicationProvider.getApplicationContext())
+                new VirtualMachineConfig.Builder(getContext())
                         .setProtectedVm(mProtectedVm)
                         .setPayloadBinaryName("MicrodroidTestNativeLib.so")
                         .setDebugLevel(debuggable ? DEBUG_LEVEL_FULL : DEBUG_LEVEL_NONE)
@@ -1537,14 +1504,7 @@
                         .build();
         final VirtualMachine vm = forceCreateNewVirtualMachine("test_vm_logcat", vmConfig);
 
-        VmEventListener listener =
-                new VmEventListener() {
-                    @Override
-                    public void onPayloadStarted(VirtualMachine vm) {
-                        forceStop(vm);
-                    }
-                };
-        listener.runToFinish(TAG, vm);
+        runVmTestService(vm, (service, results) -> {});
 
         // only check logs printed after this test
         Process logcatProcess =
@@ -1829,7 +1789,7 @@
     }
 
     private File getVmFile(String vmName, String fileName) {
-        Context context = ApplicationProvider.getApplicationContext();
+        Context context = getContext();
         Path filePath = Paths.get(context.getDataDir().getPath(), "vm", vmName, fileName);
         return filePath.toFile();
     }
@@ -1865,6 +1825,7 @@
     }
 
     static class TestResults {
+
         Exception mException;
         Integer mAddInteger;
         String mAppRunProp;
@@ -1874,6 +1835,15 @@
         String mEncryptedStoragePath;
         String[] mEffectiveCapabilities;
         String mFileContent;
+        byte[] mBcc;
+
+        void assertNoException() {
+            if (mException != null) {
+                // Rethrow, wrapped in a new exception, so we get stack traces of the original
+                // failure as well as the body of the test.
+                throw new RuntimeException(mException);
+            }
+        }
     }
 
     private TestResults runVmTestService(VirtualMachine vm, RunTestsAgainstTestService testsToRun)
@@ -1910,7 +1880,7 @@
                         }
                     }
 
-                    private void quitVMService(VirtualMachine vm) {
+                    private void quitVMService() {
                         try {
                             mTestService.quit();
                         } catch (Exception e) {
@@ -1923,7 +1893,7 @@
                         Log.i(TAG, "onPayloadReady");
                         payloadReady.complete(true);
                         testVMService(vm);
-                        quitVMService(vm);
+                        quitVMService();
                     }
 
                     @Override