[test] Move runVmAttestationService method to device test base

This allows us to reuse the method in MicrodroidTests later for
VM attestation CTS tests.

Bug: 329652894
Test: atest AvfRkpdVmAttestationTestApp
Change-Id: Icd2eef6bfb2d358b0134dcc8700f8b7794c34fdf
diff --git a/tests/helper/Android.bp b/tests/helper/Android.bp
index 614c70c..41d1ba2 100644
--- a/tests/helper/Android.bp
+++ b/tests/helper/Android.bp
@@ -16,6 +16,7 @@
         "androidx.test.runner",
         "androidx.test.ext.junit",
         "com.android.microdroid.testservice-java",
+        "com.android.virt.vm_attestation.testservice-java",
         "MicrodroidTestHelper",
         "truth",
     ],
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 2c92f04..58cb45a 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
@@ -43,6 +43,8 @@
 import com.android.microdroid.test.common.DeviceProperties;
 import com.android.microdroid.test.common.MetricsProcessor;
 import com.android.microdroid.testservice.ITestService;
+import com.android.virt.vm_attestation.testservice.IAttestationService;
+import com.android.virt.vm_attestation.testservice.IAttestationService.SigningResult;
 
 import java.io.BufferedReader;
 import java.io.ByteArrayOutputStream;
@@ -523,6 +525,40 @@
         }
     }
 
+    protected SigningResult runVmAttestationService(
+            String logTag, VirtualMachine vm, byte[] challenge, byte[] messageToSign)
+            throws Exception {
+
+        CompletableFuture<Exception> exception = new CompletableFuture<>();
+        CompletableFuture<Boolean> payloadReady = new CompletableFuture<>();
+        CompletableFuture<SigningResult> signingResultFuture = new CompletableFuture<>();
+        VmEventListener listener =
+                new VmEventListener() {
+                    @Override
+                    public void onPayloadReady(VirtualMachine vm) {
+                        payloadReady.complete(true);
+                        try {
+                            IAttestationService service =
+                                    IAttestationService.Stub.asInterface(
+                                            vm.connectToVsockServer(IAttestationService.PORT));
+                            signingResultFuture.complete(
+                                    service.signWithAttestationKey(challenge, messageToSign));
+                        } catch (Exception e) {
+                            exception.complete(e);
+                        } finally {
+                            forceStop(vm);
+                        }
+                    }
+                };
+        listener.runToFinish(TAG, vm);
+
+        assertThat(payloadReady.getNow(false)).isTrue();
+        assertThat(exception.getNow(null)).isNull();
+        SigningResult signingResult = signingResultFuture.getNow(null);
+        assertThat(signingResult).isNotNull();
+        return signingResult;
+    }
+
     protected TestResults runVmTestService(
             String logTag, VirtualMachine vm, RunTestsAgainstTestService testsToRun)
             throws Exception {