diff --git a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
index 91fb44e..ca91296 100644
--- a/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
+++ b/src/com/android/launcher3/model/AddWorkspaceItemsTask.java
@@ -15,22 +15,17 @@
  */
 package com.android.launcher3.model;
 
-import static com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID;
-
 import android.content.Intent;
 import android.content.pm.LauncherActivityInfo;
 import android.content.pm.LauncherApps;
 import android.content.pm.PackageInstaller.SessionInfo;
 import android.os.UserHandle;
 import android.util.Log;
-import android.util.LongSparseArray;
 import android.util.Pair;
 
-import com.android.launcher3.InvariantDeviceProfile;
 import com.android.launcher3.LauncherAppState;
 import com.android.launcher3.LauncherModel.CallbackTask;
 import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.config.FeatureFlags;
 import com.android.launcher3.logging.FileLog;
 import com.android.launcher3.model.BgDataModel.Callbacks;
 import com.android.launcher3.model.data.AppInfo;
@@ -41,9 +36,7 @@
 import com.android.launcher3.pm.InstallSessionHelper;
 import com.android.launcher3.pm.PackageInstallInfo;
 import com.android.launcher3.testing.TestProtocol;
-import com.android.launcher3.util.GridOccupancy;
 import com.android.launcher3.util.IntArray;
-import com.android.launcher3.util.IntSet;
 import com.android.launcher3.util.PackageManagerHelper;
 
 import java.util.ArrayList;
@@ -58,11 +51,23 @@
 
     private final List<Pair<ItemInfo, Object>> mItemList;
 
+    private final WorkspaceItemSpaceFinder mItemSpaceFinder;
+
     /**
      * @param itemList items to add on the workspace
      */
     public AddWorkspaceItemsTask(List<Pair<ItemInfo, Object>> itemList) {
+        this(itemList, new WorkspaceItemSpaceFinder());
+    }
+
+    /**
+     * @param itemList items to add on the workspace
+     * @param itemSpaceFinder inject WorkspaceItemSpaceFinder dependency for testing
+     */
+    public AddWorkspaceItemsTask(List<Pair<ItemInfo, Object>> itemList,
+            WorkspaceItemSpaceFinder itemSpaceFinder) {
         mItemList = itemList;
+        mItemSpaceFinder = itemSpaceFinder;
     }
 
     @Override
@@ -74,7 +79,7 @@
         final ArrayList<ItemInfo> addedItemsFinal = new ArrayList<>();
         final IntArray addedWorkspaceScreensFinal = new IntArray();
 
-        synchronized(dataModel) {
+        synchronized (dataModel) {
             IntArray workspaceScreens = dataModel.collectWorkspaceScreens();
 
             List<ItemInfo> filteredItems = new ArrayList<>();
@@ -117,7 +122,7 @@
 
             for (ItemInfo item : filteredItems) {
                 // Find appropriate space for the item.
-                int[] coords = findSpaceForItem(app, dataModel, workspaceScreens,
+                int[] coords = mItemSpaceFinder.findSpaceForItem(app, dataModel, workspaceScreens,
                         addedWorkspaceScreensFinal, item.spanX, item.spanY);
                 int screenId = coords[0];
 
@@ -288,82 +293,4 @@
         }
         return false;
     }
-
-    /**
-     * Find a position on the screen for the given size or adds a new screen.
-     * @return screenId and the coordinates for the item in an int array of size 3.
-     */
-    protected int[] findSpaceForItem( LauncherAppState app, BgDataModel dataModel,
-            IntArray workspaceScreens, IntArray addedWorkspaceScreensFinal, int spanX, int spanY) {
-        LongSparseArray<ArrayList<ItemInfo>> screenItems = new LongSparseArray<>();
-
-        // Use sBgItemsIdMap as all the items are already loaded.
-        synchronized (dataModel) {
-            for (ItemInfo info : dataModel.itemsIdMap) {
-                if (info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
-                    ArrayList<ItemInfo> items = screenItems.get(info.screenId);
-                    if (items == null) {
-                        items = new ArrayList<>();
-                        screenItems.put(info.screenId, items);
-                    }
-                    items.add(info);
-                }
-            }
-        }
-
-        // Find appropriate space for the item.
-        int screenId = 0;
-        int[] coordinates = new int[2];
-        boolean found = false;
-
-        int screenCount = workspaceScreens.size();
-        // First check the preferred screen.
-        IntSet screensToExclude = new IntSet();
-        if (FeatureFlags.QSB_ON_FIRST_SCREEN) {
-            screensToExclude.add(FIRST_SCREEN_ID);
-        }
-
-        for (int screen = 0; screen < screenCount; screen++) {
-            screenId = workspaceScreens.get(screen);
-            if (!screensToExclude.contains(screenId) && findNextAvailableIconSpaceInScreen(
-                    app, screenItems.get(screenId), coordinates, spanX, spanY)) {
-                // We found a space for it
-                found = true;
-                break;
-            }
-        }
-
-        if (!found) {
-            // Still no position found. Add a new screen to the end.
-            screenId = LauncherSettings.Settings.call(app.getContext().getContentResolver(),
-                    LauncherSettings.Settings.METHOD_NEW_SCREEN_ID)
-                    .getInt(LauncherSettings.Settings.EXTRA_VALUE);
-
-            // Save the screen id for binding in the workspace
-            workspaceScreens.add(screenId);
-            addedWorkspaceScreensFinal.add(screenId);
-
-            // If we still can't find an empty space, then God help us all!!!
-            if (!findNextAvailableIconSpaceInScreen(
-                    app, screenItems.get(screenId), coordinates, spanX, spanY)) {
-                throw new RuntimeException("Can't find space to add the item");
-            }
-        }
-        return new int[] {screenId, coordinates[0], coordinates[1]};
-    }
-
-    private boolean findNextAvailableIconSpaceInScreen(
-            LauncherAppState app, ArrayList<ItemInfo> occupiedPos,
-            int[] xy, int spanX, int spanY) {
-        InvariantDeviceProfile profile = app.getInvariantDeviceProfile();
-
-        GridOccupancy occupied = new GridOccupancy(profile.numColumns, profile.numRows);
-        if (occupiedPos != null) {
-            for (ItemInfo r : occupiedPos) {
-                occupied.markCells(r, true);
-            }
-        }
-        return occupied.findVacantCell(xy, spanX, spanY);
-    }
-
 }
diff --git a/src/com/android/launcher3/model/WorkspaceItemSpaceFinder.java b/src/com/android/launcher3/model/WorkspaceItemSpaceFinder.java
new file mode 100644
index 0000000..93fc6a5
--- /dev/null
+++ b/src/com/android/launcher3/model/WorkspaceItemSpaceFinder.java
@@ -0,0 +1,115 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.model;
+
+import static com.android.launcher3.WorkspaceLayoutManager.FIRST_SCREEN_ID;
+
+import android.util.LongSparseArray;
+
+import com.android.launcher3.InvariantDeviceProfile;
+import com.android.launcher3.LauncherAppState;
+import com.android.launcher3.LauncherSettings;
+import com.android.launcher3.config.FeatureFlags;
+import com.android.launcher3.model.data.ItemInfo;
+import com.android.launcher3.util.GridOccupancy;
+import com.android.launcher3.util.IntArray;
+import com.android.launcher3.util.IntSet;
+
+import java.util.ArrayList;
+
+/**
+ * Utility class to help find space for new workspace items
+ */
+public class WorkspaceItemSpaceFinder {
+
+    /**
+     * Find a position on the screen for the given size or adds a new screen.
+     *
+     * @return screenId and the coordinates for the item in an int array of size 3.
+     */
+    public int[] findSpaceForItem(LauncherAppState app, BgDataModel dataModel,
+            IntArray workspaceScreens, IntArray addedWorkspaceScreensFinal, int spanX, int spanY) {
+        LongSparseArray<ArrayList<ItemInfo>> screenItems = new LongSparseArray<>();
+
+        // Use sBgItemsIdMap as all the items are already loaded.
+        synchronized (dataModel) {
+            for (ItemInfo info : dataModel.itemsIdMap) {
+                if (info.container == LauncherSettings.Favorites.CONTAINER_DESKTOP) {
+                    ArrayList<ItemInfo> items = screenItems.get(info.screenId);
+                    if (items == null) {
+                        items = new ArrayList<>();
+                        screenItems.put(info.screenId, items);
+                    }
+                    items.add(info);
+                }
+            }
+        }
+
+        // Find appropriate space for the item.
+        int screenId = 0;
+        int[] coordinates = new int[2];
+        boolean found = false;
+
+        int screenCount = workspaceScreens.size();
+        // First check the preferred screen.
+        IntSet screensToExclude = new IntSet();
+        if (FeatureFlags.QSB_ON_FIRST_SCREEN) {
+            screensToExclude.add(FIRST_SCREEN_ID);
+        }
+
+        for (int screen = 0; screen < screenCount; screen++) {
+            screenId = workspaceScreens.get(screen);
+            if (!screensToExclude.contains(screenId) && findNextAvailableIconSpaceInScreen(
+                    app, screenItems.get(screenId), coordinates, spanX, spanY)) {
+                // We found a space for it
+                found = true;
+                break;
+            }
+        }
+
+        if (!found) {
+            // Still no position found. Add a new screen to the end.
+            screenId = LauncherSettings.Settings.call(app.getContext().getContentResolver(),
+                    LauncherSettings.Settings.METHOD_NEW_SCREEN_ID)
+                    .getInt(LauncherSettings.Settings.EXTRA_VALUE);
+
+            // Save the screen id for binding in the workspace
+            workspaceScreens.add(screenId);
+            addedWorkspaceScreensFinal.add(screenId);
+
+            // If we still can't find an empty space, then God help us all!!!
+            if (!findNextAvailableIconSpaceInScreen(
+                    app, screenItems.get(screenId), coordinates, spanX, spanY)) {
+                throw new RuntimeException("Can't find space to add the item");
+            }
+        }
+        return new int[]{screenId, coordinates[0], coordinates[1]};
+    }
+
+    private boolean findNextAvailableIconSpaceInScreen(
+            LauncherAppState app, ArrayList<ItemInfo> occupiedPos,
+            int[] xy, int spanX, int spanY) {
+        InvariantDeviceProfile profile = app.getInvariantDeviceProfile();
+
+        GridOccupancy occupied = new GridOccupancy(profile.numColumns, profile.numRows);
+        if (occupiedPos != null) {
+            for (ItemInfo r : occupiedPos) {
+                occupied.markCells(r, true);
+            }
+        }
+        return occupied.findVacantCell(xy, spanX, spanY);
+    }
+}
diff --git a/tests/src/com/android/launcher3/model/AbstractWorkspaceModelTest.kt b/tests/src/com/android/launcher3/model/AbstractWorkspaceModelTest.kt
new file mode 100644
index 0000000..d26381d
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/AbstractWorkspaceModelTest.kt
@@ -0,0 +1,162 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.model
+
+import android.content.ComponentName
+import android.content.Context
+import android.content.Intent
+import android.graphics.Rect
+import com.android.launcher3.InvariantDeviceProfile
+import com.android.launcher3.LauncherAppState
+import com.android.launcher3.LauncherSettings
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.util.ContentWriter
+import com.android.launcher3.util.GridOccupancy
+import com.android.launcher3.util.IntArray
+import com.android.launcher3.util.IntSparseArrayMap
+import com.android.launcher3.util.LauncherModelHelper
+import java.util.UUID
+
+/**
+ * Base class for workspace related tests.
+ */
+abstract class AbstractWorkspaceModelTest {
+    companion object {
+        val emptyScreenSpaces = listOf(Rect(0, 0, 5, 5))
+        val fullScreenSpaces = emptyList<Rect>()
+        val nonEmptyScreenSpaces = listOf(Rect(1, 2, 3, 4))
+    }
+
+    protected lateinit var mTargetContext: Context
+    protected lateinit var mIdp: InvariantDeviceProfile
+    protected lateinit var mAppState: LauncherAppState
+    protected lateinit var mModelHelper: LauncherModelHelper
+    protected lateinit var mExistingScreens: IntArray
+    protected lateinit var mNewScreens: IntArray
+    protected lateinit var mScreenOccupancy: IntSparseArrayMap<GridOccupancy>
+
+    open fun setup() {
+        mModelHelper = LauncherModelHelper()
+        mTargetContext = mModelHelper.sandboxContext
+        mIdp = InvariantDeviceProfile.INSTANCE[mTargetContext]
+        mIdp.numRows = 5
+        mIdp.numColumns = mIdp.numRows
+        mAppState = LauncherAppState.getInstance(mTargetContext)
+        mExistingScreens = IntArray()
+        mScreenOccupancy = IntSparseArrayMap()
+        mNewScreens = IntArray()
+    }
+
+    open fun tearDown() {
+        mModelHelper.destroy()
+    }
+
+
+    /**
+     * Sets up workspaces with the given screen IDs with some items and a 2x2 space.
+     */
+    fun setupWorkspaces(screenIdsWithItems: List<Int>) {
+        var nextItemId = 1
+        screenIdsWithItems.forEach { screenId ->
+            nextItemId = setupWorkspace(nextItemId, screenId, nonEmptyScreenSpaces)
+        }
+    }
+
+    /**
+     * Sets up the given workspaces with the given spaces, and fills the remaining space with items.
+     */
+    fun setupWorkspacesWithSpaces(
+        screen0: List<Rect>? = null,
+        screen1: List<Rect>? = null,
+        screen2: List<Rect>? = null,
+        screen3: List<Rect>? = null,
+    ) = listOf(screen0, screen1, screen2, screen3)
+        .let(this::setupWithSpaces)
+
+    private fun setupWithSpaces(workspaceSpaces: List<List<Rect>?>) {
+        var nextItemId = 1
+        workspaceSpaces.forEachIndexed { screenId, spaces ->
+            if (spaces != null) {
+                nextItemId = setupWorkspace(nextItemId, screenId, spaces)
+            }
+        }
+    }
+
+    private fun setupWorkspace(startId: Int, screenId: Int, spaces: List<Rect>): Int {
+        return mModelHelper.executeSimpleTask { dataModel ->
+            writeWorkspaceWithSpaces(dataModel, startId, screenId, spaces)
+        }
+    }
+
+    private fun writeWorkspaceWithSpaces(
+        bgDataModel: BgDataModel,
+        itemStartId: Int,
+        screenId: Int,
+        spaces: List<Rect>,
+    ): Int {
+        var itemId = itemStartId
+        val occupancy = GridOccupancy(mIdp.numColumns, mIdp.numRows)
+        occupancy.markCells(0, 0, mIdp.numColumns, mIdp.numRows, true)
+        spaces.forEach { spaceRect ->
+            occupancy.markCells(spaceRect, false)
+        }
+        mExistingScreens.add(screenId)
+        mScreenOccupancy.append(screenId, occupancy)
+        for (x in 0 until mIdp.numColumns) {
+            for (y in 0 until mIdp.numRows) {
+                if (!occupancy.cells[x][y]) {
+                    continue
+                }
+                val info = getExistingItem()
+                info.id = itemId++
+                info.screenId = screenId
+                info.cellX = x
+                info.cellY = y
+                info.container = LauncherSettings.Favorites.CONTAINER_DESKTOP
+                bgDataModel.addItem(mTargetContext, info, false)
+                val writer = ContentWriter(mTargetContext)
+                info.writeToValues(writer)
+                writer.put(LauncherSettings.Favorites._ID, info.id)
+                mTargetContext.contentResolver.insert(
+                    LauncherSettings.Favorites.CONTENT_URI,
+                    writer.getValues(mTargetContext)
+                )
+            }
+        }
+        return itemId
+    }
+
+    fun getExistingItem() = WorkspaceItemInfo()
+        .apply { intent = Intent().setComponent(ComponentName("a", "b")) }
+
+    fun getNewItem(): WorkspaceItemInfo {
+        val itemPackage = UUID.randomUUID().toString()
+        return WorkspaceItemInfo()
+            .apply { intent = Intent().setComponent(ComponentName(itemPackage, itemPackage)) }
+    }
+}
+
+data class NewItemSpace(
+    val screenId: Int,
+    val cellX: Int,
+    val cellY: Int
+) {
+    fun toIntArray() = intArrayOf(screenId, cellX, cellY)
+
+    companion object {
+        fun fromIntArray(array: kotlin.IntArray) = NewItemSpace(array[0], array[1], array[2])
+    }
+}
\ No newline at end of file
diff --git a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
deleted file mode 100644
index 8a4590a..0000000
--- a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.java
+++ /dev/null
@@ -1,201 +0,0 @@
-package com.android.launcher3.model;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.any;
-import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
-
-import android.content.ComponentName;
-import android.content.Context;
-import android.content.Intent;
-import android.graphics.Rect;
-import android.util.Pair;
-
-import androidx.test.ext.junit.runners.AndroidJUnit4;
-import androidx.test.filters.SmallTest;
-
-import com.android.launcher3.InvariantDeviceProfile;
-import com.android.launcher3.LauncherAppState;
-import com.android.launcher3.LauncherSettings;
-import com.android.launcher3.LauncherSettings.Favorites;
-import com.android.launcher3.model.BgDataModel.Callbacks;
-import com.android.launcher3.model.data.ItemInfo;
-import com.android.launcher3.model.data.WorkspaceItemInfo;
-import com.android.launcher3.util.ContentWriter;
-import com.android.launcher3.util.Executors;
-import com.android.launcher3.util.GridOccupancy;
-import com.android.launcher3.util.IntArray;
-import com.android.launcher3.util.IntSparseArrayMap;
-import com.android.launcher3.util.LauncherModelHelper;
-
-import org.junit.After;
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.ArgumentCaptor;
-
-import java.util.ArrayList;
-import java.util.List;
-
-/**
- * Tests for {@link AddWorkspaceItemsTask}
- */
-@SmallTest
-@RunWith(AndroidJUnit4.class)
-public class AddWorkspaceItemsTaskTest {
-
-    private final ComponentName mComponent1 = new ComponentName("a", "b");
-    private final ComponentName mComponent2 = new ComponentName("b", "b");
-
-    private Context mTargetContext;
-    private InvariantDeviceProfile mIdp;
-    private LauncherAppState mAppState;
-    private LauncherModelHelper mModelHelper;
-
-    private IntArray mExistingScreens;
-    private IntArray mNewScreens;
-    private IntSparseArrayMap<GridOccupancy> mScreenOccupancy;
-
-    @Before
-    public void setup() {
-        mModelHelper = new LauncherModelHelper();
-        mTargetContext = mModelHelper.sandboxContext;
-        mIdp = InvariantDeviceProfile.INSTANCE.get(mTargetContext);
-        mIdp.numColumns = mIdp.numRows = 5;
-        mAppState = LauncherAppState.getInstance(mTargetContext);
-
-        mExistingScreens = new IntArray();
-        mScreenOccupancy = new IntSparseArrayMap<>();
-        mNewScreens = new IntArray();
-    }
-
-    @After
-    public void tearDown() {
-        mModelHelper.destroy();
-    }
-
-    private AddWorkspaceItemsTask newTask(ItemInfo... items) {
-        List<Pair<ItemInfo, Object>> list = new ArrayList<>();
-        for (ItemInfo item : items) {
-            list.add(Pair.create(item, null));
-        }
-        return new AddWorkspaceItemsTask(list);
-    }
-
-    @Test
-    public void testFindSpaceForItem_prefers_second() throws Exception {
-        // First screen has only one hole of size 1
-        int nextId = setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
-
-        // Second screen has 2 holes of sizes 3x2 and 2x3
-        setupWorkspaceWithHoles(nextId, 2, new Rect(2, 0, 5, 2), new Rect(0, 2, 2, 5));
-
-        int[] spaceFound = newTask().findSpaceForItem(
-                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 1, 1);
-        assertEquals(1, spaceFound[0]);
-        assertTrue(mScreenOccupancy.get(spaceFound[0])
-                .isRegionVacant(spaceFound[1], spaceFound[2], 1, 1));
-
-        // Find a larger space
-        spaceFound = newTask().findSpaceForItem(
-                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 2, 3);
-        assertEquals(2, spaceFound[0]);
-        assertTrue(mScreenOccupancy.get(spaceFound[0])
-                .isRegionVacant(spaceFound[1], spaceFound[2], 2, 3));
-    }
-
-    @Test
-    public void testFindSpaceForItem_adds_new_screen() throws Exception {
-        // First screen has 2 holes of sizes 3x2 and 2x3
-        setupWorkspaceWithHoles(1, 1, new Rect(2, 0, 5, 2), new Rect(0, 2, 2, 5));
-
-        IntArray oldScreens = mExistingScreens.clone();
-        int[] spaceFound = newTask().findSpaceForItem(
-                mAppState, mModelHelper.getBgDataModel(), mExistingScreens, mNewScreens, 3, 3);
-        assertFalse(oldScreens.contains(spaceFound[0]));
-        assertTrue(mNewScreens.contains(spaceFound[0]));
-    }
-
-    @Test
-    public void testAddItem_existing_item_ignored() throws Exception {
-        WorkspaceItemInfo info = new WorkspaceItemInfo();
-        info.intent = new Intent().setComponent(mComponent1);
-
-        // Setup a screen with a hole
-        setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
-
-        // Nothing was added
-        assertTrue(mModelHelper.executeTaskForTest(newTask(info)).isEmpty());
-    }
-
-    @Test
-    public void testAddItem_some_items_added() throws Exception {
-        Callbacks callbacks = mock(Callbacks.class);
-        Executors.MAIN_EXECUTOR.submit(() -> mModelHelper.getModel().addCallbacks(callbacks)).get();
-
-        WorkspaceItemInfo info = new WorkspaceItemInfo();
-        info.intent = new Intent().setComponent(mComponent1);
-
-        WorkspaceItemInfo info2 = new WorkspaceItemInfo();
-        info2.intent = new Intent().setComponent(mComponent2);
-
-        // Setup a screen with a hole
-        setupWorkspaceWithHoles(1, 1, new Rect(2, 2, 3, 3));
-
-        mModelHelper.executeTaskForTest(newTask(info, info2)).get(0).run();
-        ArgumentCaptor<ArrayList> notAnimated = ArgumentCaptor.forClass(ArrayList.class);
-        ArgumentCaptor<ArrayList> animated = ArgumentCaptor.forClass(ArrayList.class);
-
-        // only info2 should be added because info was already added to the workspace
-        // in setupWorkspaceWithHoles()
-        verify(callbacks).bindAppsAdded(any(IntArray.class), notAnimated.capture(),
-                animated.capture());
-        assertTrue(notAnimated.getValue().isEmpty());
-
-        assertEquals(1, animated.getValue().size());
-        assertTrue(animated.getValue().contains(info2));
-    }
-
-    private int setupWorkspaceWithHoles(int startId, int screenId, Rect... holes) throws Exception {
-        return mModelHelper.executeSimpleTask(
-                model -> writeWorkspaceWithHoles(model, startId, screenId, holes));
-    }
-
-    private int writeWorkspaceWithHoles(
-            BgDataModel bgDataModel, int startId, int screenId, Rect... holes) {
-        GridOccupancy occupancy = new GridOccupancy(mIdp.numColumns, mIdp.numRows);
-        occupancy.markCells(0, 0, mIdp.numColumns, mIdp.numRows, true);
-        for (Rect r : holes) {
-            occupancy.markCells(r, false);
-        }
-
-        mExistingScreens.add(screenId);
-        mScreenOccupancy.append(screenId, occupancy);
-
-        for (int x = 0; x < mIdp.numColumns; x++) {
-            for (int y = 0; y < mIdp.numRows; y++) {
-                if (!occupancy.cells[x][y]) {
-                    continue;
-                }
-
-                WorkspaceItemInfo info = new WorkspaceItemInfo();
-                info.intent = new Intent().setComponent(mComponent1);
-                info.id = startId++;
-                info.screenId = screenId;
-                info.cellX = x;
-                info.cellY = y;
-                info.container = LauncherSettings.Favorites.CONTAINER_DESKTOP;
-                bgDataModel.addItem(mTargetContext, info, false);
-
-                ContentWriter writer = new ContentWriter(mTargetContext);
-                info.writeToValues(writer);
-                writer.put(Favorites._ID, info.id);
-                mTargetContext.getContentResolver().insert(Favorites.CONTENT_URI,
-                        writer.getValues(mTargetContext));
-            }
-        }
-        return startId;
-    }
-}
diff --git a/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt
new file mode 100644
index 0000000..65d938b
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/AddWorkspaceItemsTaskTest.kt
@@ -0,0 +1,245 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.model
+
+import android.util.Pair
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.launcher3.model.data.ItemInfo
+import com.android.launcher3.model.data.WorkspaceItemInfo
+import com.android.launcher3.util.Executors
+import com.android.launcher3.util.IntArray
+import com.android.launcher3.util.same
+import com.android.launcher3.util.eq
+import com.android.launcher3.util.any
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.ArgumentCaptor
+import org.mockito.Captor
+import org.mockito.Mock
+import org.mockito.MockitoAnnotations
+import org.mockito.Mockito.verify
+import org.mockito.Mockito.verifyZeroInteractions
+import org.mockito.Mockito.times
+import org.mockito.Mockito.`when` as whenever
+
+/**
+ * Tests for [AddWorkspaceItemsTask]
+ */
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class AddWorkspaceItemsTaskTest : AbstractWorkspaceModelTest() {
+
+    @Captor
+    private lateinit var mAnimatedItemArgumentCaptor: ArgumentCaptor<ArrayList<ItemInfo>>
+
+    @Captor
+    private lateinit var mNotAnimatedItemArgumentCaptor: ArgumentCaptor<ArrayList<ItemInfo>>
+
+    @Mock
+    private lateinit var mDataModelCallbacks: BgDataModel.Callbacks
+
+    @Mock
+    private lateinit var mWorkspaceItemSpaceFinder: WorkspaceItemSpaceFinder
+
+
+    @Before
+    override fun setup() {
+        super.setup()
+        MockitoAnnotations.initMocks(this)
+        Executors.MAIN_EXECUTOR.submit { mModelHelper.model.addCallbacks(mDataModelCallbacks) }
+            .get()
+    }
+
+    @After
+    override fun tearDown() {
+        super.tearDown()
+    }
+
+    @Test
+    fun givenNewItemAndNonEmptyPages_whenExecuteTask_thenAddNewItem() {
+        val itemToAdd = getNewItem()
+        val nonEmptyScreenIds = listOf(0, 1, 2)
+        givenNewItemSpaces(NewItemSpace(1, 2, 2))
+
+        val addedItems = testAddItems(nonEmptyScreenIds, itemToAdd)
+
+        assertThat(addedItems.size).isEqualTo(1)
+        assertThat(addedItems.first().itemInfo.screenId).isEqualTo(1)
+        assertThat(addedItems.first().isAnimated).isTrue()
+        verifyItemSpaceFinderCall(nonEmptyScreenIds, numberOfExpectedCall = 1)
+    }
+
+    @Test
+    fun givenNewAndExistingItems_whenExecuteTask_thenOnlyAddNewItem() {
+        val itemsToAdd = arrayOf(
+            getNewItem(),
+            getExistingItem()
+        )
+        givenNewItemSpaces(NewItemSpace(1, 0, 0))
+        val nonEmptyScreenIds = listOf(0)
+
+        val addedItems = testAddItems(nonEmptyScreenIds, *itemsToAdd)
+
+        assertThat(addedItems.size).isEqualTo(1)
+        assertThat(addedItems.first().itemInfo.screenId).isEqualTo(1)
+        assertThat(addedItems.first().isAnimated).isTrue()
+        verifyItemSpaceFinderCall(nonEmptyScreenIds, numberOfExpectedCall = 1)
+    }
+
+    @Test
+    fun givenOnlyExistingItem_whenExecuteTask_thenDoNotAddItem() {
+        val itemToAdd = getExistingItem()
+        givenNewItemSpaces(NewItemSpace(1, 0, 0))
+        val nonEmptyScreenIds = listOf(0)
+
+        val addedItems = testAddItems(nonEmptyScreenIds, itemToAdd)
+
+        assertThat(addedItems.size).isEqualTo(0)
+        verifyZeroInteractions(mWorkspaceItemSpaceFinder, mDataModelCallbacks)
+    }
+
+    @Test
+    fun givenNonSequentialScreenIds_whenExecuteTask_thenReturnNewScreenId() {
+        val itemToAdd = getNewItem()
+        givenNewItemSpaces(NewItemSpace(2, 1, 3))
+        val nonEmptyScreenIds = listOf(0, 2, 3)
+
+        val addedItems = testAddItems(nonEmptyScreenIds, itemToAdd)
+
+        assertThat(addedItems.size).isEqualTo(1)
+        assertThat(addedItems.first().itemInfo.screenId).isEqualTo(2)
+        assertThat(addedItems.first().isAnimated).isTrue()
+        verifyItemSpaceFinderCall(nonEmptyScreenIds, numberOfExpectedCall = 1)
+    }
+
+    @Test
+    fun givenMultipleItems_whenExecuteTask_thenAddThem() {
+        val itemsToAdd = arrayOf(
+            getNewItem(),
+            getExistingItem(),
+            getNewItem(),
+            getNewItem(),
+            getExistingItem(),
+        )
+        givenNewItemSpaces(
+            NewItemSpace(1, 3, 3),
+            NewItemSpace(2, 0, 0),
+            NewItemSpace(2, 0, 1),
+        )
+        val nonEmptyScreenIds = listOf(0, 1)
+
+        val addedItems = testAddItems(nonEmptyScreenIds, *itemsToAdd)
+
+        // Only the new items should be added
+        assertThat(addedItems.size).isEqualTo(3)
+
+        // Items that are added to the first screen should not be animated
+        val itemsAddedToFirstScreen = addedItems.filter { it.itemInfo.screenId == 1 }
+        assertThat(itemsAddedToFirstScreen.size).isEqualTo(1)
+        assertThat(itemsAddedToFirstScreen.first().isAnimated).isFalse()
+
+        // Items that are added to the second screen should be animated
+        val itemsAddedToSecondScreen = addedItems.filter { it.itemInfo.screenId == 2 }
+        assertThat(itemsAddedToSecondScreen.size).isEqualTo(2)
+        itemsAddedToSecondScreen.forEach {
+            assertThat(it.isAnimated).isTrue()
+        }
+        verifyItemSpaceFinderCall(nonEmptyScreenIds, numberOfExpectedCall = 3)
+    }
+
+    /**
+     * Sets up the item space data that will be returned from WorkspaceItemSpaceFinder.
+     */
+    private fun givenNewItemSpaces(vararg newItemSpaces: NewItemSpace) {
+        val spaceStack = newItemSpaces.toMutableList()
+        whenever(
+            mWorkspaceItemSpaceFinder.findSpaceForItem(
+                any(),
+                any(),
+                any(),
+                any(),
+                any(),
+                any()
+            )
+        )
+            .then { spaceStack.removeFirst().toIntArray() }
+    }
+
+    /**
+     * Verifies if WorkspaceItemSpaceFinder was called with proper arguments and how many times was
+     * it called.
+     */
+    private fun verifyItemSpaceFinderCall(
+        nonEmptyScreenIds: List<Int>,
+        numberOfExpectedCall: Int
+    ) {
+        verify(mWorkspaceItemSpaceFinder, times(numberOfExpectedCall))
+            .findSpaceForItem(
+                same(mAppState), same(mModelHelper.bgDataModel),
+                eq(IntArray.wrap(*nonEmptyScreenIds.toIntArray())), eq(IntArray()), eq(1), eq(1)
+            )
+    }
+
+    /**
+     * Sets up the workspaces with items, executes the task, collects the added items from the
+     * model callback then returns it.
+     */
+    private fun testAddItems(
+        nonEmptyScreenIds: List<Int>,
+        vararg itemsToAdd: WorkspaceItemInfo
+    ): List<AddedItem> {
+        setupWorkspaces(nonEmptyScreenIds)
+        val task = newTask(*itemsToAdd)
+        var updateCount = 0
+        mModelHelper.executeTaskForTest(task)
+            .forEach {
+                updateCount++
+                it.run()
+            }
+
+        val addedItems = mutableListOf<AddedItem>()
+        if (updateCount > 0) {
+            verify(mDataModelCallbacks).bindAppsAdded(
+                any(),
+                mNotAnimatedItemArgumentCaptor.capture(), mAnimatedItemArgumentCaptor.capture()
+            )
+            addedItems.addAll(mAnimatedItemArgumentCaptor.value.map { AddedItem(it, true) })
+            addedItems.addAll(mNotAnimatedItemArgumentCaptor.value.map { AddedItem(it, false) })
+
+        }
+
+        return addedItems
+    }
+
+    /**
+     * Creates the task with the given items and replaces the WorkspaceItemSpaceFinder dependency
+     * with a mock.
+     */
+    private fun newTask(vararg items: ItemInfo): AddWorkspaceItemsTask =
+        items.map { Pair.create(it, Any()) }
+            .toMutableList()
+            .let { AddWorkspaceItemsTask(it, mWorkspaceItemSpaceFinder) }
+}
+
+private data class AddedItem(
+    val itemInfo: ItemInfo,
+    val isAnimated: Boolean
+)
diff --git a/tests/src/com/android/launcher3/model/WorkspaceItemSpaceFinderTest.kt b/tests/src/com/android/launcher3/model/WorkspaceItemSpaceFinderTest.kt
new file mode 100644
index 0000000..bfb1ac6
--- /dev/null
+++ b/tests/src/com/android/launcher3/model/WorkspaceItemSpaceFinderTest.kt
@@ -0,0 +1,171 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.launcher3.model
+
+import android.graphics.Rect
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth.assertThat
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+/**
+ * Tests for [WorkspaceItemSpaceFinder]
+ */
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class WorkspaceItemSpaceFinderTest : AbstractWorkspaceModelTest() {
+
+    private val mItemSpaceFinder = WorkspaceItemSpaceFinder()
+
+    @Before
+    override fun setup() {
+        super.setup()
+    }
+
+    @After
+    override fun tearDown() {
+        super.tearDown()
+    }
+
+    private fun findSpace(spanX: Int, spanY: Int): NewItemSpace =
+        mItemSpaceFinder.findSpaceForItem(
+            mAppState, mModelHelper.bgDataModel,
+            mExistingScreens, mNewScreens, spanX, spanY
+        )
+            .let { NewItemSpace.fromIntArray(it) }
+
+    private fun assertRegionVacant(newItemSpace: NewItemSpace, spanX: Int, spanY: Int) {
+        assertThat(
+            mScreenOccupancy[newItemSpace.screenId]
+                .isRegionVacant(newItemSpace.cellX, newItemSpace.cellY, spanX, spanY)
+        ).isTrue()
+    }
+
+    @Test
+    fun justEnoughSpaceOnFirstScreen_whenFindSpaceForItem_thenReturnFirstScreenId() {
+        setupWorkspacesWithSpaces(
+            // 3x2 space on screen 0, but it should be skipped
+            screen0 = listOf(Rect(2, 0, 5, 2)),
+            screen1 = listOf(Rect(2, 2, 3, 3)), // 1x1 space
+            //  2 spaces of sizes 3x2 and 2x3
+            screen2 = listOf(Rect(2, 0, 5, 2), Rect(0, 2, 2, 5)),
+        )
+
+        val spaceFound = findSpace(1, 1)
+
+        assertThat(spaceFound.screenId).isEqualTo(1)
+        assertRegionVacant(spaceFound, 1, 1)
+    }
+
+    @Test
+    fun notEnoughSpaceOnFirstScreen_whenFindSpaceForItem_thenReturnSecondScreenId() {
+        setupWorkspacesWithSpaces(
+            // 3x2 space on screen 0, but it should be skipped
+            screen0 = listOf(Rect(2, 0, 5, 2)),
+            screen1 = listOf(Rect(2, 2, 3, 3)), // 1x1 space
+            //  2 spaces of sizes 3x2 and 2x3
+            screen2 = listOf(Rect(2, 0, 5, 2), Rect(0, 2, 2, 5)),
+        )
+
+        // Find a larger space
+        val spaceFound = findSpace(2, 3)
+
+        assertThat(spaceFound.screenId).isEqualTo(2)
+        assertRegionVacant(spaceFound, 2, 3)
+    }
+
+    @Test
+    fun notEnoughSpaceOnExistingScreens_returnNewScreenId() {
+        setupWorkspacesWithSpaces(
+            // 3x2 space on screen 0, but it should be skipped
+            screen0 = listOf(Rect(2, 0, 5, 2)),
+            //  2 spaces of sizes 3x2 and 2x3
+            screen1 = listOf(Rect(2, 0, 5, 2), Rect(0, 2, 2, 5)),
+            //  2 spaces of sizes 1x2 and 2x2
+            screen2 = listOf(Rect(1, 0, 2, 2), Rect(3, 2, 5, 4)),
+        )
+
+        val oldScreens = mExistingScreens.clone()
+        val spaceFound = findSpace(3, 3)
+
+        assertThat(oldScreens.contains(spaceFound.screenId)).isFalse()
+        assertThat(mNewScreens.contains(spaceFound.screenId)).isTrue()
+    }
+
+    @Test
+    fun firstScreenIsEmptyButSecondIsNotEmpty_returnSecondScreenId() {
+        setupWorkspacesWithSpaces(
+            // 3x2 space on screen 0, but it should be skipped
+            screen0 = listOf(Rect(2, 0, 5, 2)),
+            // empty screens are skipped
+            screen2 = listOf(Rect(2, 0, 5, 2)), // 3x2 space
+        )
+
+        val spaceFound = findSpace(2, 1)
+
+        assertThat(spaceFound.screenId).isEqualTo(2)
+        assertRegionVacant(spaceFound, 2, 1)
+    }
+
+    @Test
+    fun twoEmptyMiddleScreens_returnThirdScreen() {
+        setupWorkspacesWithSpaces(
+            // 3x2 space on screen 0, but it should be skipped
+            screen0 = listOf(Rect(2, 0, 5, 2)),
+            // empty screens are skipped
+            screen3 = listOf(Rect(1, 1, 4, 4)), // 3x3 space
+        )
+
+        val spaceFound = findSpace(2, 3)
+
+        assertThat(spaceFound.screenId).isEqualTo(3)
+        assertRegionVacant(spaceFound, 2, 3)
+    }
+
+    @Test
+    fun allExistingPagesAreFull_returnNewScreenId() {
+        setupWorkspacesWithSpaces(
+            // 3x2 space on screen 0, but it should be skipped
+            screen0 = listOf(Rect(2, 0, 5, 2)),
+            screen1 = fullScreenSpaces,
+            screen2 = fullScreenSpaces,
+        )
+
+        val spaceFound = findSpace(2, 3)
+
+        assertThat(spaceFound.screenId).isEqualTo(3)
+        assertThat(mNewScreens.contains(spaceFound.screenId)).isTrue()
+    }
+
+    @Test
+    fun firstTwoPagesAreFull_and_ThirdPageIsEmpty_returnThirdPage() {
+        setupWorkspacesWithSpaces(
+            // 3x2 space on screen 0, but it should be skipped
+            screen0 = listOf(Rect(2, 0, 5, 2)),
+            screen1 = fullScreenSpaces, // full screens are skipped
+            screen2 = fullScreenSpaces, // full screens are skipped
+            screen3 = emptyScreenSpaces
+        )
+
+        val spaceFound = findSpace(3, 1)
+
+        assertThat(spaceFound.screenId).isEqualTo(3)
+        assertRegionVacant(spaceFound, 3, 1)
+    }
+}
diff --git a/tests/src/com/android/launcher3/util/KotlinMockitoHelpers.kt b/tests/src/com/android/launcher3/util/KotlinMockitoHelpers.kt
new file mode 100644
index 0000000..57db13a
--- /dev/null
+++ b/tests/src/com/android/launcher3/util/KotlinMockitoHelpers.kt
@@ -0,0 +1,117 @@
+/*
+ * Copyright (C) 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.launcher3.util
+
+/**
+ * Kotlin versions of popular mockito methods that can return null in situations when Kotlin expects
+ * a non-null value. Kotlin will throw an IllegalStateException when this takes place ("x must not
+ * be null"). To fix this, we can use methods that modify the return type to be nullable. This
+ * causes Kotlin to skip the null checks.
+ */
+
+import org.mockito.ArgumentCaptor
+import org.mockito.Mockito
+
+/**
+ * Returns Mockito.eq() as nullable type to avoid java.lang.IllegalStateException when
+ * null is returned.
+ *
+ * Generic T is nullable because implicitly bounded by Any?.
+ */
+fun <T> eq(obj: T): T = Mockito.eq<T>(obj)
+
+/**
+ * Returns Mockito.same() as nullable type to avoid java.lang.IllegalStateException when
+ * null is returned.
+ *
+ * Generic T is nullable because implicitly bounded by Any?.
+ */
+fun <T> same(obj: T): T = Mockito.same<T>(obj)
+
+/**
+ * Returns Mockito.any() as nullable type to avoid java.lang.IllegalStateException when
+ * null is returned.
+ *
+ * Generic T is nullable because implicitly bounded by Any?.
+ */
+fun <T> any(type: Class<T>): T = Mockito.any<T>(type)
+inline fun <reified T> any(): T = any(T::class.java)
+
+/**
+ * Kotlin type-inferred version of Mockito.nullable()
+ */
+inline fun <reified T> nullable(): T? = Mockito.nullable(T::class.java)
+
+/**
+ * Returns ArgumentCaptor.capture() as nullable type to avoid java.lang.IllegalStateException
+ * when null is returned.
+ *
+ * Generic T is nullable because implicitly bounded by Any?.
+ */
+fun <T> capture(argumentCaptor: ArgumentCaptor<T>): T = argumentCaptor.capture()
+
+/**
+ * Helper function for creating an argumentCaptor in kotlin.
+ *
+ * Generic T is nullable because implicitly bounded by Any?.
+ */
+inline fun <reified T : Any> argumentCaptor(): ArgumentCaptor<T> =
+    ArgumentCaptor.forClass(T::class.java)
+
+/**
+ * Helper function for creating new mocks, without the need to pass in a [Class] instance.
+ *
+ * Generic T is nullable because implicitly bounded by Any?.
+ */
+inline fun <reified T : Any> mock(): T = Mockito.mock(T::class.java)
+
+/**
+ * A kotlin implemented wrapper of [ArgumentCaptor] which prevents the following exception when
+ * kotlin tests are mocking kotlin objects and the methods take non-null parameters:
+ *
+ *     java.lang.NullPointerException: capture() must not be null
+ */
+class KotlinArgumentCaptor<T> constructor(clazz: Class<T>) {
+    private val wrapped: ArgumentCaptor<T> = ArgumentCaptor.forClass(clazz)
+    fun capture(): T = wrapped.capture()
+    val value: T
+        get() = wrapped.value
+}
+
+/**
+ * Helper function for creating an argumentCaptor in kotlin.
+ *
+ * Generic T is nullable because implicitly bounded by Any?.
+ */
+inline fun <reified T : Any> kotlinArgumentCaptor(): KotlinArgumentCaptor<T> =
+    KotlinArgumentCaptor(T::class.java)
+
+/**
+ * Helper function for creating and using a single-use ArgumentCaptor in kotlin.
+ *
+ *    val captor = argumentCaptor<Foo>()
+ *    verify(...).someMethod(captor.capture())
+ *    val captured = captor.value
+ *
+ * becomes:
+ *
+ *    val captured = withArgCaptor<Foo> { verify(...).someMethod(capture()) }
+ *
+ * NOTE: this uses the KotlinArgumentCaptor to avoid the NullPointerException.
+ */
+inline fun <reified T : Any> withArgCaptor(block: KotlinArgumentCaptor<T>.() -> Unit): T =
+    kotlinArgumentCaptor<T>().apply { block() }.value