diff --git a/service_vm/test_apk/Android.bp b/service_vm/test_apk/Android.bp
index 450d475..72e411e 100644
--- a/service_vm/test_apk/Android.bp
+++ b/service_vm/test_apk/Android.bp
@@ -58,14 +58,8 @@
     manifest: "AndroidManifest.rkpd.xml",
     test_config: "AndroidTest.rkpd.xml",
     static_libs: [
-        "RkpdAppTestUtil",
         "VmAttestationTestUtil",
-        "androidx.work_work-testing",
     ],
-    instrumentation_for: "rkpdapp",
-    // This app is a variation of rkpdapp, with additional permissions to run
-    // a VM. It is defined in packages/modules/RemoteKeyProvisioning.
-    data: [":avf-rkpdapp"],
 }
 
 java_library {
diff --git a/service_vm/test_apk/AndroidManifest.rkpd.xml b/service_vm/test_apk/AndroidManifest.rkpd.xml
index 6ecc5a9..369a456 100644
--- a/service_vm/test_apk/AndroidManifest.rkpd.xml
+++ b/service_vm/test_apk/AndroidManifest.rkpd.xml
@@ -17,8 +17,15 @@
 <manifest xmlns:android="http://schemas.android.com/apk/res/android"
     package="com.android.virt.rkpd.vm_attestation.testapp">
 
+    <uses-permission android:name="android.permission.MANAGE_VIRTUAL_MACHINE" />
+    <uses-permission android:name="android.permission.USE_CUSTOM_VIRTUAL_MACHINE" />
+
+    <!-- Required to check the network access -->
+    <uses-permission android:name="android.permission.ACCESS_NETWORK_STATE" />
+    <uses-permission android:name="android.permission.INTERNET" />
+
     <instrumentation
             android:name="androidx.test.runner.AndroidJUnitRunner"
-            android:targetPackage="com.android.rkpdapp"
+            android:targetPackage="com.android.virt.rkpd.vm_attestation.testapp"
             android:label="AVF rkpd app integration tests" />
 </manifest>
diff --git a/service_vm/test_apk/AndroidTest.rkpd.xml b/service_vm/test_apk/AndroidTest.rkpd.xml
index 39eca32..8ff7bb4 100644
--- a/service_vm/test_apk/AndroidTest.rkpd.xml
+++ b/service_vm/test_apk/AndroidTest.rkpd.xml
@@ -17,22 +17,18 @@
     <option name="test-suite-tag" value="apct" />
     <option name="test-suite-tag" value="apct-instrumentation" />
 
-    <!-- Need to disable SELinux policy to allow com.android.rkpdapp to run a VM. -->
-    <target_preparer class="com.android.tradefed.targetprep.DisableSELinuxTargetPreparer"/>
+    <target_preparer class="com.android.tradefed.targetprep.DeviceSetup">
+        <!-- The host name is needed for RKPD key provisioning -->
+        <option name="set-property" key="remote_provisioning.hostname"
+         value="remoteprovisioning.googleapis.com" />
+        <option name="restore-properties" value="true"/>
+    </target_preparer>
 
     <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
         <option name="test-file-name" value="AvfRkpdVmAttestationTestApp.apk" />
-        <option name="test-file-name" value="avf-rkpdapp.apk" />
     </target_preparer>
 
     <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
         <option name="package" value="com.android.virt.rkpd.vm_attestation.testapp" />
     </test>
-
-    <!-- Only run if RKPD mainline module is installed -->
-    <object type="module_controller"
-            class="com.android.tradefed.testtype.suite.module.MainlineTestModuleController">
-        <option name="enable" value="true" />
-        <option name="mainline-module-package-name" value="com.android.rkpd" />
-    </object>
 </configuration>
diff --git a/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java b/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java
index ce7fc45..678e56f 100644
--- a/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java
+++ b/service_vm/test_apk/src/java/com/android/virt/rkpd/vm_attestation/testapp/RkpdVmAttestationTest.java
@@ -18,31 +18,20 @@
 
 import static android.system.virtualmachine.VirtualMachineConfig.DEBUG_LEVEL_FULL;
 
-import static com.google.common.truth.Truth.assertThat;
 import static com.google.common.truth.TruthJUnit.assume;
 
+import android.net.ConnectivityManager;
+import android.net.NetworkCapabilities;
+import android.net.Network;
 import android.content.Context;
-import android.hardware.security.keymint.IRemotelyProvisionedComponent;
-import android.os.SystemProperties;
 import android.system.virtualmachine.VirtualMachine;
 import android.system.virtualmachine.VirtualMachineConfig;
 
-import androidx.work.ListenableWorker;
-import androidx.work.testing.TestWorkerBuilder;
-
 import com.android.microdroid.test.device.MicrodroidDeviceTestBase;
-import com.android.rkpdapp.database.ProvisionedKeyDao;
-import com.android.rkpdapp.database.RkpdDatabase;
-import com.android.rkpdapp.interfaces.ServerInterface;
-import com.android.rkpdapp.interfaces.ServiceManagerInterface;
-import com.android.rkpdapp.interfaces.SystemInterface;
-import com.android.rkpdapp.provisioner.PeriodicProvisioner;
-import com.android.rkpdapp.testutil.SystemInterfaceSelector;
-import com.android.rkpdapp.utils.Settings;
 import com.android.virt.vm_attestation.testservice.IAttestationService.SigningResult;
 import com.android.virt.vm_attestation.util.X509Utils;
+import android.system.virtualmachine.VirtualMachineManager;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -53,7 +42,6 @@
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
-import java.util.concurrent.Executors;
 
 /**
  * End-to-end test for the pVM remote attestation.
@@ -71,23 +59,16 @@
  * <p>- Have an arm64 device supporting protected VMs.
  *
  * <p>- Have a stable network connection on the device.
- *
- * <p>- Have the RKP server hostname configured in the device. If not, you can set it using: $ adb
- * shell setprop remote_provisioning.hostname remoteprovisioning.googleapis.com
  */
 @RunWith(Parameterized.class)
 public class RkpdVmAttestationTest extends MicrodroidDeviceTestBase {
     private static final String TAG = "RkpdVmAttestationTest";
 
-    private static final String SERVICE_NAME = IRemotelyProvisionedComponent.DESCRIPTOR + "/avf";
     private static final String VM_PAYLOAD_PATH = "libvm_attestation_test_payload.so";
     private static final String MESSAGE = "Hello RKP from AVF!";
     private static final String TEST_APP_PACKAGE_NAME =
             "com.android.virt.rkpd.vm_attestation.testapp";
 
-    private ProvisionedKeyDao mKeyDao;
-    private PeriodicProvisioner mProvisioner;
-
     @Parameterized.Parameter(0)
     public String mGki;
 
@@ -103,57 +84,28 @@
 
     @Before
     public void setUp() throws Exception {
-        assume().withMessage("The RKP server hostname is not configured -- assume RKP disabled.")
-                .that(SystemProperties.get("remote_provisioning.hostname"))
-                .isNotEmpty();
         assume().withMessage("RKP Integration tests rely on network availability.")
-                .that(ServerInterface.isNetworkConnected(getContext()))
+                .that(isNetworkConnected(getContext()))
                 .isTrue();
-        // TODO(b/329652894): Assume that pVM remote attestation feature is supported.
+        assumeFeatureEnabled(VirtualMachineManager.FEATURE_REMOTE_ATTESTATION);
+        assume().withMessage("Test needs Remote Attestation support")
+                .that(getVirtualMachineManager().isRemoteAttestationSupported())
+                .isTrue();
 
-        prepareTestSetup(true /* protectedVm */, mGki);
-
-        Settings.clearPreferences(getContext());
-        mKeyDao = RkpdDatabase.getDatabase(getContext()).provisionedKeyDao();
-        mKeyDao.deleteAllKeys();
-
-        mProvisioner =
-                TestWorkerBuilder.from(
-                                getContext(),
-                                PeriodicProvisioner.class,
-                                Executors.newSingleThreadExecutor())
-                        .build();
-
-        SystemInterface systemInterface =
-                SystemInterfaceSelector.getSystemInterfaceForServiceName(SERVICE_NAME);
-        ServiceManagerInterface.setInstances(new SystemInterface[] {systemInterface});
-
-        setMaxPerformanceTaskProfile();
-    }
-
-    @After
-    public void tearDown() throws Exception {
-        ServiceManagerInterface.setInstances(null);
-        if (mKeyDao != null) {
-            mKeyDao.deleteAllKeys();
+        if (mGki == null) {
+            // We don't need this permission to use the microdroid kernel.
+            revokePermission(VirtualMachine.USE_CUSTOM_VIRTUAL_MACHINE_PERMISSION);
         }
-        Settings.clearPreferences(getContext());
+        prepareTestSetup(true /* protectedVm */, mGki);
+        setMaxPerformanceTaskProfile();
     }
 
     @Test
     public void usingProvisionedKeyForVmAttestationSucceeds() throws Exception {
-        // Provision keys.
-        assertThat(mProvisioner.doWork()).isEqualTo(ListenableWorker.Result.success());
-        assertThat(mKeyDao.getTotalUnassignedKeysForIrpc(SERVICE_NAME)).isGreaterThan(0);
-
         // Arrange.
-        Context ctx = getContext();
-        Context otherAppCtx = ctx.createPackageContext(TEST_APP_PACKAGE_NAME, 0);
         VirtualMachineConfig config =
-                new VirtualMachineConfig.Builder(otherAppCtx)
-                        .setProtectedVm(true)
+                newVmConfigBuilderWithPayloadBinary(VM_PAYLOAD_PATH)
                         .setDebugLevel(DEBUG_LEVEL_FULL)
-                        .setPayloadBinaryName(VM_PAYLOAD_PATH)
                         .setVmOutputCaptured(true)
                         .build();
         VirtualMachine vm = forceCreateNewVirtualMachine("attestation_with_rkpd_client", config);
@@ -170,4 +122,13 @@
         X509Utils.verifyAvfRelatedCerts(certs, challenge, TEST_APP_PACKAGE_NAME);
         X509Utils.verifySignature(certs[0], MESSAGE.getBytes(), signingResult.signature);
     }
+
+    private static boolean isNetworkConnected(Context context) {
+        ConnectivityManager cm = context.getSystemService(ConnectivityManager.class);
+        Network network = cm.getActiveNetwork();
+        NetworkCapabilities capabilities = cm.getNetworkCapabilities(network);
+        return capabilities != null
+                && capabilities.hasCapability(NetworkCapabilities.NET_CAPABILITY_INTERNET)
+                && capabilities.hasCapability(NetworkCapabilities.NET_CAPABILITY_VALIDATED);
+    }
 }