diff --git a/cmds/am/src/com/android/commands/am/Am.java b/cmds/am/src/com/android/commands/am/Am.java
index bdb8380..846a34e 100644
--- a/cmds/am/src/com/android/commands/am/Am.java
+++ b/cmds/am/src/com/android/commands/am/Am.java
@@ -174,10 +174,6 @@
                 instrument.noWindowAnimation = true;
             } else if (opt.equals("--no-hidden-api-checks")) {
                 instrument.disableHiddenApiChecks = true;
-            } else if (opt.equals("--no-test-api-checks")) {
-                // TODO(satayev): remove this option, only kept for backwards compatibility with
-                // cached tradefed instance
-                instrument.disableTestApiChecks = false;
             } else if (opt.equals("--no-test-api-access")) {
                 instrument.disableTestApiChecks = false;
             } else if (opt.equals("--no-isolated-storage")) {
@@ -198,7 +194,6 @@
         }
 
         instrument.componentNameArg = nextArgRequired();
-
         instrument.run();
     }
 }
diff --git a/services/core/java/com/android/server/am/ActivityManagerService.java b/services/core/java/com/android/server/am/ActivityManagerService.java
index 928ddab..686adbb 100644
--- a/services/core/java/com/android/server/am/ActivityManagerService.java
+++ b/services/core/java/com/android/server/am/ActivityManagerService.java
@@ -190,6 +190,7 @@
 import android.app.usage.UsageStatsManager;
 import android.app.usage.UsageStatsManagerInternal;
 import android.appwidget.AppWidgetManager;
+import android.compat.Compatibility;
 import android.content.AutofillOptions;
 import android.content.BroadcastReceiver;
 import android.content.ComponentCallbacks2;
@@ -318,6 +319,7 @@
 import com.android.internal.app.ProcessMap;
 import com.android.internal.app.SystemUserHomeActivity;
 import com.android.internal.app.procstats.ProcessStats;
+import com.android.internal.compat.CompatibilityChangeConfig;
 import com.android.internal.content.PackageHelper;
 import com.android.internal.messages.nano.SystemMessageProto.SystemMessage;
 import com.android.internal.notification.SystemNotificationChannels;
@@ -16944,6 +16946,8 @@
             if (disableHiddenApiChecks || disableTestApiChecks) {
                 enforceCallingPermission(android.Manifest.permission.DISABLE_HIDDEN_API_CHECKS,
                         "disable hidden API checks");
+
+                enableTestApiAccess(ii.packageName);
             }
 
             // TODO(b/158750470): remove
@@ -17083,6 +17087,25 @@
 
         forceStopPackageLocked(app.info.packageName, -1, false, false, true, true, false, app.userId,
                 "finished inst");
+
+        disableTestApiAccess(app.info.packageName);
+    }
+
+    private void enableTestApiAccess(String packageName) {
+        if (mPlatformCompat != null) {
+            Compatibility.ChangeConfig config = new Compatibility.ChangeConfig(
+                    Collections.singleton(166236554L /* VMRuntime.ALLOW_TEST_API_ACCESS */),
+                    Collections.emptySet());
+            CompatibilityChangeConfig override = new CompatibilityChangeConfig(config);
+            mPlatformCompat.setOverridesForTest(override, packageName);
+        }
+    }
+
+    private void disableTestApiAccess(String packageName) {
+        if (mPlatformCompat != null) {
+            mPlatformCompat.clearOverrideForTest(166236554L /* VMRuntime.ALLOW_TEST_API_ACCESS */,
+                    packageName);
+        }
     }
 
     public void finishInstrumentation(IApplicationThread target,