From 3c354421c83b2e02293ac2db0f0dd2e922400b29 Mon Sep 17 00:00:00 2001
From: Sunny Goyal <sunnygoyal@google.com>
Date: Fri, 13 May 2022 09:01:14 -0700
Subject: [PATCH] Extra checks to find the source of flaky tests

Bug: 200572078
Test: Presubmit
Change-Id: Ie927f476e3a2c1b5aee84e33bf510bb635e59071
---
 src/com/android/launcher3/model/BaseLoaderResults.java | 5 ++++-
 src/com/android/launcher3/model/BgDataModel.java       | 2 ++
 src/com/android/launcher3/model/ModelUtils.java        | 2 +-
 3 files changed, 7 insertions(+), 2 deletions(-)

diff --git a/src/com/android/launcher3/model/BaseLoaderResults.java b/src/com/android/launcher3/model/BaseLoaderResults.java
index 6c4cfb9068..b50ab587f4 100644
--- a/src/com/android/launcher3/model/BaseLoaderResults.java
+++ b/src/com/android/launcher3/model/BaseLoaderResults.java
@@ -43,6 +43,7 @@ import com.android.launcher3.util.RunnableList;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
+import java.util.Objects;
 import java.util.concurrent.Executor;
 
 /**
@@ -203,7 +204,9 @@ public abstract class BaseLoaderResults {
         }
 
         private void bind() {
-            IntSet currentScreenIds = mCallbacks.getPagesToBindSynchronously(mOrderedScreenIds);
+            final IntSet currentScreenIds =
+                    mCallbacks.getPagesToBindSynchronously(mOrderedScreenIds);
+            Objects.requireNonNull(currentScreenIds, "Null screen ids provided by " + mCallbacks);
 
             // Separate the items that are on the current screen, and all the other remaining items
             ArrayList<ItemInfo> currentWorkspaceItems = new ArrayList<>();
diff --git a/src/com/android/launcher3/model/BgDataModel.java b/src/com/android/launcher3/model/BgDataModel.java
index d52537e074..de23c4b31f 100644
--- a/src/com/android/launcher3/model/BgDataModel.java
+++ b/src/com/android/launcher3/model/BgDataModel.java
@@ -31,6 +31,7 @@ import android.text.TextUtils;
 import android.util.ArraySet;
 import android.util.Log;
 
+import androidx.annotation.NonNull;
 import androidx.annotation.Nullable;
 
 import com.android.launcher3.LauncherSettings;
@@ -469,6 +470,7 @@ public class BgDataModel {
          * or an empty IntSet
          * @param orderedScreenIds All the page ids to be bound
          */
+        @NonNull
         default IntSet getPagesToBindSynchronously(IntArray orderedScreenIds) {
             return new IntSet();
         }
diff --git a/src/com/android/launcher3/model/ModelUtils.java b/src/com/android/launcher3/model/ModelUtils.java
index df6768d582..422af43e19 100644
--- a/src/com/android/launcher3/model/ModelUtils.java
+++ b/src/com/android/launcher3/model/ModelUtils.java
@@ -51,7 +51,7 @@ public class ModelUtils {
      * specified screen.
      */
     public static <T extends ItemInfo> void filterCurrentWorkspaceItems(
-            IntSet currentScreenIds,
+            final IntSet currentScreenIds,
             ArrayList<T> allWorkspaceItems,
             ArrayList<T> currentScreenItems,
             ArrayList<T> otherScreenItems) {