Adds implementation for AccessibilityService#getWindowsOnAllDisplays to ShadowAccessibilityService.

PiperOrigin-RevId: 507044048
diff --git a/robolectric/src/test/java/org/robolectric/shadows/ShadowAccessibilityServiceTest.java b/robolectric/src/test/java/org/robolectric/shadows/ShadowAccessibilityServiceTest.java
index 9cabcdf..195d673 100644
--- a/robolectric/src/test/java/org/robolectric/shadows/ShadowAccessibilityServiceTest.java
+++ b/robolectric/src/test/java/org/robolectric/shadows/ShadowAccessibilityServiceTest.java
@@ -224,12 +224,100 @@
     w2.setId(2);
     AccessibilityWindowInfo w3 = AccessibilityWindowInfo.obtain();
     w3.setId(3);
+
     shadow.setWindows(Arrays.asList(w1, w2, w3));
+
     assertThat(service.getWindows()).hasSize(3);
     assertThat(service.getWindows()).containsExactly(w1, w2, w3).inOrder();
   }
 
   @Test
+  @Config(minSdk = R)
+  public void getWindowsforDefaultDisplay_returnEmptyList() {
+    assertThat(service.getWindowsOnAllDisplays().get(Display.DEFAULT_DISPLAY)).isEmpty();
+  }
+
+  @Test
+  @Config(minSdk = R)
+  public void getWindowsforNonDefaultDisplay_returnNullList() {
+    assertThat(service.getWindowsOnAllDisplays().get(Display.DEFAULT_DISPLAY + 1)).isNull();
+  }
+
+  @Test
+  @Config(minSdk = R)
+  public void setWindowsOnDisplay_returnPopulatedWindowsOnAllDisplays() {
+    AccessibilityWindowInfo w1 = AccessibilityWindowInfo.obtain();
+    w1.setId(1);
+    AccessibilityWindowInfo w2 = AccessibilityWindowInfo.obtain();
+    w2.setId(2);
+    AccessibilityWindowInfo w3 = AccessibilityWindowInfo.obtain();
+    w3.setId(3);
+    AccessibilityWindowInfo w4 = AccessibilityWindowInfo.obtain();
+    w4.setId(4);
+    AccessibilityWindowInfo w5 = AccessibilityWindowInfo.obtain();
+    w5.setId(5);
+    AccessibilityWindowInfo w6 = AccessibilityWindowInfo.obtain();
+    w6.setId(6);
+    AccessibilityWindowInfo w7 = AccessibilityWindowInfo.obtain();
+    w7.setId(7);
+
+    shadow.setWindowsOnDisplay(Display.DEFAULT_DISPLAY, Arrays.asList(w1, w2, w3));
+    shadow.setWindowsOnDisplay(Display.DEFAULT_DISPLAY + 1, Arrays.asList(w4, w5, w6, w7));
+
+    assertThat(service.getWindows()).hasSize(3);
+    assertThat(service.getWindows()).containsExactly(w1, w2, w3).inOrder();
+    assertThat(service.getWindowsOnAllDisplays().size()).isEqualTo(2);
+    assertThat(service.getWindowsOnAllDisplays().get(Display.DEFAULT_DISPLAY))
+        .containsExactly(w1, w2, w3)
+        .inOrder();
+    assertThat(service.getWindowsOnAllDisplays().get(Display.DEFAULT_DISPLAY + 1))
+        .containsExactly(w4, w5, w6, w7)
+        .inOrder();
+  }
+
+  @Test
+  @Config(minSdk = R)
+  public void setNullWindowsOnNonDefaultDisplay_nonDefaultDisplayHasWindows_displayIsRemoved() {
+    AccessibilityWindowInfo w1 = AccessibilityWindowInfo.obtain();
+    w1.setId(1);
+    AccessibilityWindowInfo w2 = AccessibilityWindowInfo.obtain();
+    w2.setId(2);
+    AccessibilityWindowInfo w3 = AccessibilityWindowInfo.obtain();
+    w3.setId(3);
+    shadow.setWindowsOnDisplay(Display.DEFAULT_DISPLAY + 1, Arrays.asList(w1, w2, w3));
+
+    shadow.setWindowsOnDisplay(Display.DEFAULT_DISPLAY + 1, null);
+
+    assertThat(service.getWindowsOnAllDisplays().get(Display.DEFAULT_DISPLAY + 1)).isNull();
+  }
+
+  @Test
+  @Config(minSdk = R)
+  public void setWindows_nonDefaultDisplayHasWindows_nonDefaultDisplayWindowsNotRemoved() {
+    AccessibilityWindowInfo w1 = AccessibilityWindowInfo.obtain();
+    w1.setId(1);
+    AccessibilityWindowInfo w2 = AccessibilityWindowInfo.obtain();
+    w2.setId(2);
+    AccessibilityWindowInfo w3 = AccessibilityWindowInfo.obtain();
+    w3.setId(3);
+    AccessibilityWindowInfo w4 = AccessibilityWindowInfo.obtain();
+    w4.setId(4);
+    AccessibilityWindowInfo w5 = AccessibilityWindowInfo.obtain();
+    w5.setId(5);
+    AccessibilityWindowInfo w6 = AccessibilityWindowInfo.obtain();
+    w6.setId(6);
+    AccessibilityWindowInfo w7 = AccessibilityWindowInfo.obtain();
+    w7.setId(7);
+    shadow.setWindowsOnDisplay(Display.DEFAULT_DISPLAY + 1, Arrays.asList(w4, w5, w6, w7));
+
+    shadow.setWindows(Arrays.asList(w1, w2, w3));
+
+    assertThat(service.getWindowsOnAllDisplays().get(Display.DEFAULT_DISPLAY + 1))
+        .containsExactly(w4, w5, w6, w7)
+        .inOrder();
+  }
+
+  @Test
   @Config(minSdk = S)
   public void getSystemActions_returnsNull() {
     assertThat(service.getSystemActions()).isNull();
diff --git a/shadows/framework/src/main/java/org/robolectric/shadows/ShadowAccessibilityService.java b/shadows/framework/src/main/java/org/robolectric/shadows/ShadowAccessibilityService.java
index e5620df..9e51896 100644
--- a/shadows/framework/src/main/java/org/robolectric/shadows/ShadowAccessibilityService.java
+++ b/shadows/framework/src/main/java/org/robolectric/shadows/ShadowAccessibilityService.java
@@ -16,11 +16,15 @@
 import android.hardware.HardwareBuffer;
 import android.os.Handler;
 import android.os.SystemClock;
+import android.util.SparseArray;
+import android.view.Display;
 import android.view.accessibility.AccessibilityNodeInfo;
 import android.view.accessibility.AccessibilityWindowInfo;
+import com.google.common.collect.ArrayListMultimap;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.concurrent.Executor;
+import javax.annotation.Nullable;
 import org.robolectric.annotation.Implementation;
 import org.robolectric.annotation.Implements;
 import org.robolectric.util.ReflectionHelpers;
@@ -35,7 +39,8 @@
 
   private final List<Integer> globalActionsPerformed = new ArrayList<>();
   private List<AccessibilityNodeInfo.AccessibilityAction> systemActions;
-  private final List<AccessibilityWindowInfo> windows = new ArrayList<>();
+  private final ArrayListMultimap<Integer, AccessibilityWindowInfo> windows =
+      ArrayListMultimap.create();
   private final List<GestureDispatch> gesturesDispatched = new ArrayList<>();
 
   private boolean canDispatchGestures = true;
@@ -66,13 +71,41 @@
   }
 
   /**
-   * Returns a representation of interactive windows shown on the device screen. Mirrors the values
-   * provided to {@link #setWindows(List<AccessibilityWindowInfo>)}. Returns an empty List if not
-   * set.
+   * Returns a representation of interactive windows shown on the device's default display. Mirrors
+   * the values provided to {@link #setWindows(List<AccessibilityWindowInfo>)}. Returns an empty
+   * list if not set.
    */
   @Implementation(minSdk = LOLLIPOP)
   protected List<AccessibilityWindowInfo> getWindows() {
-    return new ArrayList<>(windows);
+    List<AccessibilityWindowInfo> windowInfos = windows.get(Display.DEFAULT_DISPLAY);
+    if (windowInfos != null) {
+      return new ArrayList<>(windowInfos);
+    } else {
+      return new ArrayList<>();
+    }
+  }
+
+  /**
+   * Returns a representation of interactive windows shown on the device's all displays. An empty
+   * list will be returned for default display and {@code null} will be return for other displays if
+   * they are not set by {@link #setWindowsOnDisplay(int, List)}.
+   */
+  @Implementation(minSdk = R)
+  protected SparseArray<List<AccessibilityWindowInfo>> getWindowsOnAllDisplays() {
+    return cloneWindowOnAllDisplays();
+  }
+
+  private SparseArray<List<AccessibilityWindowInfo>> cloneWindowOnAllDisplays() {
+    SparseArray<List<AccessibilityWindowInfo>> anotherArray = new SparseArray<>();
+    for (int displayId : windows.keySet()) {
+      anotherArray.put(displayId, windows.get(displayId));
+    }
+
+    if (anotherArray.get(Display.DEFAULT_DISPLAY) == null) {
+      anotherArray.put(Display.DEFAULT_DISPLAY, new ArrayList<>());
+    }
+
+    return anotherArray;
   }
 
   @Implementation(minSdk = N)
@@ -130,13 +163,23 @@
   }
 
   /**
-   * Sets the list of interactive windows shown on the device screen as reported by {@link
-   * #getWindows()}
+   * Sets the list of interactive windows shown on the device's default display as reported by
+   * {@link #getWindows()}
    */
   public void setWindows(List<AccessibilityWindowInfo> windowList) {
-    windows.clear();
+    setWindowsOnDisplay(Display.DEFAULT_DISPLAY, windowList);
+  }
+
+  /**
+   * Sets the list of interactive windows shown on the device's {@code displayId} display. If the
+   * {@code windowList} is null, we will remove the list with given {@code displayId} display as
+   * reported by {@link #getWindowsOnAllDisplays()}.
+   */
+  public void setWindowsOnDisplay(
+      int displayId, @Nullable List<AccessibilityWindowInfo> windowList) {
+    windows.removeAll(displayId);
     if (windowList != null) {
-      windows.addAll(windowList);
+      windows.putAll(displayId, windowList);
     }
   }