diff --git a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
index 1acecfc..a23636c 100644
--- a/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
+++ b/services/tests/servicestests/src/com/android/server/devicepolicy/DevicePolicyManagerTest.java
@@ -3726,7 +3726,7 @@
 
     private void verifyLockTaskState(int userId, String[] packages, int flags) throws Exception {
         verify(getServices().iactivityManager).updateLockTaskPackages(userId, packages);
-        verify(getServices().iactivityManager).updateLockTaskFeatures(userId, flags);
+        verify(getServices().iactivityTaskManager).updateLockTaskFeatures(userId, flags);
     }
 
     private void verifyCanSetLockTask(int uid, int userId, ComponentName who, String[] packages,
@@ -3819,7 +3819,7 @@
         // Lock task packages cleared when loading user data and when the user becomes unaffiliated.
         verify(getServices().iactivityManager, times(2)).updateLockTaskPackages(
                 MANAGED_PROFILE_USER_ID, new String[0]);
-        verify(getServices().iactivityManager, times(2)).updateLockTaskFeatures(
+        verify(getServices().iactivityTaskManager, times(2)).updateLockTaskFeatures(
                 MANAGED_PROFILE_USER_ID, DevicePolicyManager.LOCK_TASK_FEATURE_NONE);
 
         // Verify that lock task packages were not cleared for the DO