Extract runVmTestService into MicrodroidDeviceTestBase

This is a helpful method to build tests around. Extract it into the base
class so it can also be used in benchmarks.

Test: atest MicrodroidTestApp
Change-Id: I0a055d1f5c423e42fbf447812a06776bae76a966
diff --git a/tests/helper/Android.bp b/tests/helper/Android.bp
index 61c5dcd..c9eafad 100644
--- a/tests/helper/Android.bp
+++ b/tests/helper/Android.bp
@@ -15,6 +15,7 @@
     static_libs: [
         "androidx.test.runner",
         "androidx.test.ext.junit",
+        "com.android.microdroid.testservice-java",
         "MicrodroidTestHelper",
         "truth-prebuilt",
     ],
diff --git a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
index 9ec36b3..be3c1da 100644
--- a/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
+++ b/tests/helper/src/java/com/android/microdroid/test/device/MicrodroidDeviceTestBase.java
@@ -17,6 +17,7 @@
 
 import static android.content.pm.PackageManager.FEATURE_VIRTUALIZATION_FRAMEWORK;
 
+import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.TruthJUnit.assume;
 
 import android.app.Instrumentation;
@@ -38,6 +39,7 @@
 
 import com.android.microdroid.test.common.DeviceProperties;
 import com.android.microdroid.test.common.MetricsProcessor;
+import com.android.microdroid.testservice.ITestService;
 
 import java.io.BufferedReader;
 import java.io.ByteArrayOutputStream;
@@ -428,4 +430,104 @@
             throw new RuntimeException("Failed to run the command.");
         }
     }
+
+    protected static class TestResults {
+        public Exception mException;
+        public Integer mAddInteger;
+        public String mAppRunProp;
+        public String mSublibRunProp;
+        public String mExtraApkTestProp;
+        public String mApkContentsPath;
+        public String mEncryptedStoragePath;
+        public String[] mEffectiveCapabilities;
+        public String mFileContent;
+        public byte[] mBcc;
+        public long[] mTimings;
+
+        public 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);
+            }
+        }
+    }
+
+    protected TestResults runVmTestService(
+            String logTag, VirtualMachine vm, RunTestsAgainstTestService testsToRun)
+            throws Exception {
+        CompletableFuture<Boolean> payloadStarted = new CompletableFuture<>();
+        CompletableFuture<Boolean> payloadReady = new CompletableFuture<>();
+        CompletableFuture<Boolean> payloadFinished = new CompletableFuture<>();
+        TestResults testResults = new TestResults();
+        VmEventListener listener =
+                new VmEventListener() {
+                    ITestService mTestService = null;
+
+                    private void initializeTestService(VirtualMachine vm) {
+                        try {
+                            mTestService =
+                                    ITestService.Stub.asInterface(
+                                            vm.connectToVsockServer(ITestService.SERVICE_PORT));
+                            // Make sure linkToDeath works, and include it in the log in case it's
+                            // helpful.
+                            mTestService
+                                    .asBinder()
+                                    .linkToDeath(
+                                            () -> Log.i(logTag, "ITestService binder died"), 0);
+                        } catch (Exception e) {
+                            testResults.mException = e;
+                        }
+                    }
+
+                    private void testVMService(VirtualMachine vm) {
+                        try {
+                            if (mTestService == null) initializeTestService(vm);
+                            testsToRun.runTests(mTestService, testResults);
+                        } catch (Exception e) {
+                            testResults.mException = e;
+                        }
+                    }
+
+                    private void quitVMService() {
+                        try {
+                            mTestService.quit();
+                        } catch (Exception e) {
+                            testResults.mException = e;
+                        }
+                    }
+
+                    @Override
+                    public void onPayloadReady(VirtualMachine vm) {
+                        Log.i(logTag, "onPayloadReady");
+                        payloadReady.complete(true);
+                        testVMService(vm);
+                        quitVMService();
+                    }
+
+                    @Override
+                    public void onPayloadStarted(VirtualMachine vm) {
+                        Log.i(logTag, "onPayloadStarted");
+                        payloadStarted.complete(true);
+                    }
+
+                    @Override
+                    public void onPayloadFinished(VirtualMachine vm, int exitCode) {
+                        Log.i(logTag, "onPayloadFinished: " + exitCode);
+                        payloadFinished.complete(true);
+                        forceStop(vm);
+                    }
+                };
+
+        listener.runToFinish(logTag, vm);
+        assertThat(payloadStarted.getNow(false)).isTrue();
+        assertThat(payloadReady.getNow(false)).isTrue();
+        assertThat(payloadFinished.getNow(false)).isTrue();
+        return testResults;
+    }
+
+    @FunctionalInterface
+    protected interface RunTestsAgainstTestService {
+        void runTests(ITestService testService, TestResults testResults) throws Exception;
+    }
 }