diff --git a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
index 731cd1e3a503a..12375e8a6611a 100644
--- a/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
+++ b/services/devicepolicy/java/com/android/server/devicepolicy/DevicePolicyManagerService.java
@@ -2290,6 +2290,11 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
                     context, requestCode, intent, flags, options, user);
         }
 
+        PendingIntent pendingIntentGetBroadcast(
+                Context context, int requestCode, Intent intent, int flags) {
+            return PendingIntent.getBroadcast(context, requestCode, intent, flags);
+        }
+
         void registerContentObserver(Uri uri, boolean notifyForDescendents,
                 ContentObserver observer, int userHandle) {
             mContext.getContentResolver().registerContentObserver(uri, notifyForDescendents,
@@ -15666,8 +15671,8 @@ public class DevicePolicyManagerService extends BaseIDevicePolicyManager {
 
     private void updateProfileOffAlarm(long profileOffDeadline) {
         final AlarmManager am = mInjector.getAlarmManager();
-        final PendingIntent pi = PendingIntent.getBroadcast(mContext, REQUEST_PROFILE_OFF_DEADLINE,
-                new Intent(ACTION_PROFILE_OFF_DEADLINE),
+        final PendingIntent pi = mInjector.pendingIntentGetBroadcast(
+                mContext, REQUEST_PROFILE_OFF_DEADLINE, new Intent(ACTION_PROFILE_OFF_DEADLINE),
                 PendingIntent.FLAG_ONE_SHOT | PendingIntent.FLAG_UPDATE_CURRENT);
         am.cancel(pi);
         if (profileOffDeadline != 0) {
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceMigrationTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceMigrationTest.java
index 9574a086c74b5..40b0e7114cc92 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceMigrationTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceMigrationTest.java
@@ -31,8 +31,11 @@ import android.content.pm.PackageManager;
 import android.os.Bundle;
 import android.os.UserHandle;
 import android.os.UserManager;
+import android.platform.test.annotations.Presubmit;
 import android.provider.Settings;
 
+import androidx.test.filters.SmallTest;
+
 import com.android.frameworks.servicestests.R;
 import com.android.server.LocalServices;
 import com.android.server.SystemService;
@@ -43,7 +46,7 @@ import java.util.HashMap;
 import java.util.Map;
 import java.util.Set;
 
-// TODO (b/143516163): Fix old test cases and put into presubmit.
+// TODO (b/149818286): Fix old test cases and put the whole test into presubmit.
 public class DevicePolicyManagerServiceMigrationTest extends DpmTestBase {
 
     private static final String USER_TYPE_EMPTY = "";
@@ -343,6 +346,8 @@ public class DevicePolicyManagerServiceMigrationTest extends DpmTestBase {
         assertTrue(alreadySet.contains(UserManager.DISALLOW_BLUETOOTH_SHARING));
     }
 
+    @Presubmit
+    @SmallTest
     public void testCompMigrationUnAffiliated_skipped() throws Exception {
         prepareAdmin1AsDo();
         prepareAdminAnotherPackageAsPo(COPE_PROFILE_USER_ID);
@@ -354,6 +359,8 @@ public class DevicePolicyManagerServiceMigrationTest extends DpmTestBase {
         assertTrue(dpms.mOwners.hasDeviceOwner());
     }
 
+    @Presubmit
+    @SmallTest
     public void testCompMigrationAffiliated() throws Exception {
         prepareAdmin1AsDo();
         prepareAdmin1AsPo(COPE_PROFILE_USER_ID);
diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
index 3a8258be5f010..853151f8a0de6 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerServiceTestable.java
@@ -368,6 +368,12 @@ public class DevicePolicyManagerServiceTestable extends DevicePolicyManagerServi
             return null;
         }
 
+        @Override
+        PendingIntent pendingIntentGetBroadcast(Context context, int requestCode,
+                Intent intent, int flags) {
+            return null;
+        }
+
         @Override
         void registerContentObserver(Uri uri, boolean notifyForDescendents,
                 ContentObserver observer, int userHandle) {