diff --git a/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java b/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java
index db1d021..e225a59 100644
--- a/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java
+++ b/tests/benchmark/src/java/com/android/microdroid/benchmark/BenchmarkVmListener.java
@@ -20,6 +20,7 @@
 
 import android.os.RemoteException;
 import android.system.virtualmachine.VirtualMachine;
+import android.system.virtualmachine.VirtualMachineException;
 import android.util.Log;
 
 import com.android.microdroid.test.device.MicrodroidDeviceTestBase.VmEventListener;
@@ -32,6 +33,8 @@
 class BenchmarkVmListener extends VmEventListener {
     private static final String TAG = "BenchmarkVm";
 
+    private boolean mPayloadReady;
+
     interface InnerListener {
         /** This is invoked when both the payload and {@link IBenchmarkService} are ready. */
         void onPayloadReady(VirtualMachine vm, IBenchmarkService benchmarkService)
@@ -46,6 +49,7 @@
 
     @Override
     public final void onPayloadReady(VirtualMachine vm) {
+        mPayloadReady = true;
         try {
             IBenchmarkService benchmarkService =
                     IBenchmarkService.Stub.asInterface(
@@ -60,7 +64,20 @@
         forceStop(vm);
     }
 
+    @Override
+    public void onError(VirtualMachine vm, int errorCode, String message) {
+        throw new RuntimeException("Error while benchmark (" + errorCode + "): " + message);
+    }
+
     static BenchmarkVmListener create(InnerListener listener) {
         return new BenchmarkVmListener(listener);
     }
+
+    @Override
+    public void runToFinish(String logTag, VirtualMachine vm)
+            throws VirtualMachineException, InterruptedException {
+        mPayloadReady = false;
+        super.runToFinish(logTag, vm);
+        assertThat(mPayloadReady).isTrue();
+    }
 }