Localize the removal of root tasks to default TDA in tests

- Change the implementation of removeRootTaskWithActivityTypes test
  API such that the removal only happens on default TDA.
- As there could be safety critical activities running in other TDAs
- Problem with removing those activities is that they cannot be
  reliably restarted in a way where all CTS remain fine.

Flag: EXEMPT (Test only)
Bug: 205581941
Test: Ran all WM CTS modules on a device having activities in
      additional TDAs
Change-Id: Ib1710348b19ae1a81b6228f9788db067dde08a3f
diff --git a/core/java/android/app/ActivityTaskManager.java b/core/java/android/app/ActivityTaskManager.java
index c8ab260..799df1f 100644
--- a/core/java/android/app/ActivityTaskManager.java
+++ b/core/java/android/app/ActivityTaskManager.java
@@ -172,7 +172,7 @@
         }
     }
 
-    /** Removes root tasks of the activity types from the system. */
+    /** Removes root tasks of the activity types from the Default TDA of all displays. */
     @RequiresPermission(android.Manifest.permission.MANAGE_ACTIVITY_TASKS)
     public void removeRootTasksWithActivityTypes(@NonNull int[] activityTypes) {
         try {
diff --git a/services/core/java/com/android/server/wm/DisplayContent.java b/services/core/java/com/android/server/wm/DisplayContent.java
index 1efb3ef..7147c0c 100644
--- a/services/core/java/com/android/server/wm/DisplayContent.java
+++ b/services/core/java/com/android/server/wm/DisplayContent.java
@@ -6278,7 +6278,7 @@
         // by looping the children, so that we don't miss any root tasks after the children size
         // changed or reordered.
         final ArrayList<Task> rootTasks = new ArrayList<>();
-        forAllRootTasks(rootTask -> {
+        getDefaultTaskDisplayArea().forAllRootTasks(rootTask -> {
             for (int activityType : activityTypes) {
                 // Collect the root tasks that are currently being organized.
                 if (rootTask.mCreatedByOrganizer) {