Properly remove local color areas of callback

Refactored to system server to the repository pattern
added locks to the WallpaperManager

Fixes: 206632865
Test: manual
Change-Id: Icc73c18e34042ae48e338b03dc3e94fee5d6939d

4 files changed, 363 insertions, 123 deletions

diff --git a/core/java/android/app/WallpaperManager.java b/core/java/android/app/WallpaperManager.java
index 22b1c03fba50..fca4c698c49c 100644
--- a/core/java/android/app/WallpaperManager.java
+++ b/core/java/android/app/WallpaperManager.java
@@ -365,17 +365,18 @@ public class WallpaperManager {
         private int mCachedWallpaperUserId;
         private Bitmap mDefaultWallpaper;
         private Handler mMainLooperHandler;
-        private ArrayMap<RectF, ArraySet<LocalWallpaperColorConsumer>> mLocalColorAreas =
-                new ArrayMap<>();
+        private ArrayMap<LocalWallpaperColorConsumer, ArraySet<RectF>> mLocalColorCallbackAreas =
+                        new ArrayMap<>();
         private ILocalWallpaperColorConsumer mLocalColorCallback =
                 new ILocalWallpaperColorConsumer.Stub() {
                     @Override
                     public void onColorsChanged(RectF area, WallpaperColors colors) {
-                        ArraySet<LocalWallpaperColorConsumer> callbacks =
-                                mLocalColorAreas.get(area);
-                        if (callbacks == null) return;
-                        for (LocalWallpaperColorConsumer callback: callbacks) {
-                            callback.onColorsChanged(area, colors);
+                        for (LocalWallpaperColorConsumer callback :
+                                mLocalColorCallbackAreas.keySet()) {
+                            ArraySet<RectF> areas = mLocalColorCallbackAreas.get(callback);
+                            if (areas != null && areas.contains(area)) {
+                                callback.onColorsChanged(area, colors);
+                            }
                         }
                     }
                 };
@@ -420,46 +421,52 @@ public class WallpaperManager {
             }
         }
 
-        public void addOnColorsChangedListener(@NonNull LocalWallpaperColorConsumer callback,
+        public void addOnColorsChangedListener(
+                @NonNull LocalWallpaperColorConsumer callback,
                 @NonNull List<RectF> regions, int which, int userId, int displayId) {
-            for (RectF area: regions) {
-                ArraySet<LocalWallpaperColorConsumer> callbacks = mLocalColorAreas.get(area);
-                if (callbacks == null) {
-                    callbacks = new ArraySet<>();
-                    mLocalColorAreas.put(area, callbacks);
+            synchronized (this) {
+                for (RectF area : regions) {
+                    ArraySet<RectF> areas = mLocalColorCallbackAreas.get(callback);
+                    if (areas == null) {
+                        areas = new ArraySet<>();
+                        mLocalColorCallbackAreas.put(callback, areas);
+                    }
+                    areas.add(area);
+                }
+                try {
+                    // one way returns immediately
+                    mService.addOnLocalColorsChangedListener(mLocalColorCallback, regions, which,
+                            userId, displayId);
+                } catch (RemoteException e) {
+                    // Can't get colors, connection lost.
+                    Log.e(TAG, "Can't register for local color updates", e);
                 }
-                callbacks.add(callback);
-            }
-            try {
-                mService.addOnLocalColorsChangedListener(mLocalColorCallback , regions, which,
-                                                         userId, displayId);
-            } catch (RemoteException e) {
-                // Can't get colors, connection lost.
-                Log.e(TAG, "Can't register for local color updates", e);
             }
         }
 
         public void removeOnColorsChangedListener(
                 @NonNull LocalWallpaperColorConsumer callback, int which, int userId,
                 int displayId) {
-            final ArrayList<RectF> removeAreas = new ArrayList<>();
-            for (RectF area : mLocalColorAreas.keySet()) {
-                ArraySet<LocalWallpaperColorConsumer> callbacks = mLocalColorAreas.get(area);
-                if (callbacks == null) continue;
-                callbacks.remove(callback);
-                if (callbacks.size() == 0) {
-                    mLocalColorAreas.remove(area);
-                    removeAreas.add(area);
+            synchronized (this) {
+                final ArraySet<RectF> removeAreas = mLocalColorCallbackAreas.remove(callback);
+                if (removeAreas == null || removeAreas.size() == 0) {
+                    return;
                 }
-            }
-            try {
-                if (removeAreas.size() > 0) {
-                    mService.removeOnLocalColorsChangedListener(
-                            mLocalColorCallback, removeAreas, which, userId, displayId);
+                for (LocalWallpaperColorConsumer cb : mLocalColorCallbackAreas.keySet()) {
+                    ArraySet<RectF> areas = mLocalColorCallbackAreas.get(cb);
+                    if (areas != null && cb != callback) removeAreas.removeAll(areas);
+                }
+                try {
+                    if (removeAreas.size() > 0) {
+                        // one way returns immediately
+                        mService.removeOnLocalColorsChangedListener(
+                                mLocalColorCallback, new ArrayList(removeAreas), which, userId,
+                                displayId);
+                    }
+                } catch (RemoteException e) {
+                    // Can't get colors, connection lost.
+                    Log.e(TAG, "Can't unregister for local color updates", e);
                 }
-            } catch (RemoteException e) {
-                // Can't get colors, connection lost.
-                Log.e(TAG, "Can't unregister for local color updates", e);
             }
         }
 
diff --git a/services/core/java/com/android/server/wallpaper/LocalColorRepository.java b/services/core/java/com/android/server/wallpaper/LocalColorRepository.java
new file mode 100644
index 000000000000..e6265f4dbd39
--- /dev/null
+++ b/services/core/java/com/android/server/wallpaper/LocalColorRepository.java
@@ -0,0 +1,166 @@
+/*
+ * Copyright (C) 2021 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.server.wallpaper;
+
+import android.app.ILocalWallpaperColorConsumer;
+import android.graphics.RectF;
+import android.os.IBinder;
+import android.os.RemoteCallbackList;
+import android.os.RemoteException;
+import android.util.ArrayMap;
+import android.util.ArraySet;
+import android.util.SparseArray;
+
+import com.android.internal.annotations.VisibleForTesting;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.function.Consumer;
+
+/**
+ * Manages the lifecycle of local wallpaper color callbacks and their interested wallpaper regions.
+ */
+public class LocalColorRepository {
+    /**
+     * Maps local wallpaper color callbacks' binders to their interested wallpaper regions, which
+     * are stored in a map of display Ids to wallpaper regions.
+     * binder callback -> [display id: int] -> areas
+     */
+    ArrayMap<IBinder, SparseArray<ArraySet<RectF>>> mLocalColorAreas = new ArrayMap();
+    RemoteCallbackList<ILocalWallpaperColorConsumer> mCallbacks = new RemoteCallbackList();
+
+    /**
+     * Add areas to a consumer
+     * @param consumer
+     * @param areas
+     * @param displayId
+     */
+    public void addAreas(ILocalWallpaperColorConsumer consumer, List<RectF> areas, int displayId) {
+        IBinder binder = consumer.asBinder();
+        SparseArray<ArraySet<RectF>> displays = mLocalColorAreas.get(binder);
+        ArraySet<RectF> displayAreas = null;
+        if (displays == null) {
+            try {
+                consumer.asBinder().linkToDeath(() ->
+                        mLocalColorAreas.remove(consumer.asBinder()), 0);
+            } catch (RemoteException e) {
+                e.printStackTrace();
+            }
+            displays = new SparseArray<>();
+            mLocalColorAreas.put(binder, displays);
+        } else {
+            displayAreas = displays.get(displayId);
+        }
+        if (displayAreas == null) {
+            displayAreas = new ArraySet(areas);
+            displays.put(displayId, displayAreas);
+        }
+
+        for (int i = 0; i < areas.size(); i++) {
+            displayAreas.add(areas.get(i));
+        }
+        mCallbacks.register(consumer);
+    }
+
+    /**
+     * remove an area for a consumer
+     * @param consumer
+     * @param areas
+     * @param displayId
+     * @return the areas that are removed from all callbacks
+     */
+    public List<RectF> removeAreas(ILocalWallpaperColorConsumer consumer, List<RectF> areas,
+            int displayId) {
+        IBinder binder = consumer.asBinder();
+        SparseArray<ArraySet<RectF>> displays = mLocalColorAreas.get(binder);
+        ArraySet<RectF> registeredAreas = null;
+        if (displays != null) {
+            registeredAreas = displays.get(displayId);
+            if (registeredAreas == null) {
+                mCallbacks.unregister(consumer);
+            } else {
+                for (int i = 0; i < areas.size(); i++) {
+                    registeredAreas.remove(areas.get(i));
+                }
+                if (registeredAreas.size() == 0) {
+                    displays.remove(displayId);
+                }
+            }
+            if (displays.size() == 0) {
+                mLocalColorAreas.remove(binder);
+                mCallbacks.unregister(consumer);
+            }
+        } else {
+            mCallbacks.unregister(consumer);
+        }
+        ArraySet<RectF> purged = new ArraySet<>(areas);
+        for (int i = 0; i < mLocalColorAreas.size(); i++) {
+            for (int j = 0; j < mLocalColorAreas.valueAt(i).size(); j++) {
+                for (int k = 0; k < mLocalColorAreas.valueAt(i).valueAt(j).size(); k++) {
+                    purged.remove(mLocalColorAreas.valueAt(i).valueAt(j).valueAt(k));
+                }
+            }
+        }
+        return new ArrayList(purged);
+    }
+
+    /**
+     * Return the local areas by display id
+     * @param displayId
+     * @return
+     */
+    public List<RectF> getAreasByDisplayId(int displayId) {
+        ArrayList<RectF> areas = new ArrayList();
+        for (int i = 0; i < mLocalColorAreas.size(); i++) {
+            SparseArray<ArraySet<RectF>> displays = mLocalColorAreas.valueAt(i);
+            if (displays == null) continue;
+            ArraySet<RectF> displayAreas = displays.get(displayId);
+            if (displayAreas == null) continue;
+            for (int j = 0; j < displayAreas.size(); j++) {
+                areas.add(displayAreas.valueAt(j));
+            }
+        }
+        return areas;
+    }
+
+    /**
+     * invoke a callback for each area of interest
+     * @param callback
+     * @param area
+     * @param displayId
+     */
+    public void forEachCallback(Consumer<ILocalWallpaperColorConsumer> callback,
+            RectF area, int displayId) {
+        mCallbacks.broadcast(cb -> {
+            IBinder binder = cb.asBinder();
+            SparseArray<ArraySet<RectF>> displays = mLocalColorAreas.get(binder);
+            if (displays == null) return;
+            ArraySet<RectF> displayAreas = displays.get(displayId);
+            if (displayAreas != null && displayAreas.contains(area)) callback.accept(cb);
+        });
+    }
+
+    /**
+     * For testing
+     * @param callback
+     * @return if the callback is registered
+     */
+    @VisibleForTesting
+    protected boolean isCallbackAvailable(ILocalWallpaperColorConsumer callback) {
+        return mLocalColorAreas.get(callback.asBinder()) != null;
+    }
+}
diff --git a/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java b/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
index 2f353d19b5df..a7924e61cd3e 100644
--- a/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
+++ b/services/core/java/com/android/server/wallpaper/WallpaperManagerService.java
@@ -89,8 +89,6 @@ import android.service.wallpaper.IWallpaperService;
 import android.service.wallpaper.WallpaperService;
 import android.system.ErrnoException;
 import android.system.Os;
-import android.util.ArrayMap;
-import android.util.ArraySet;
 import android.util.EventLog;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -881,12 +879,7 @@ public class WallpaperManagerService extends IWallpaperManager.Stub
     private final SparseBooleanArray mUserRestorecon = new SparseBooleanArray();
     private int mCurrentUserId = UserHandle.USER_NULL;
     private boolean mInAmbientMode;
-    private ArrayMap<IBinder, ArraySet<RectF>> mLocalColorCallbackAreas =
-            new ArrayMap<>();
-    private ArrayMap<RectF, RemoteCallbackList<ILocalWallpaperColorConsumer>>
-            mLocalColorAreaCallbacks = new ArrayMap<>();
-    private ArrayMap<Integer, ArraySet<RectF>> mLocalColorDisplayIdAreas = new ArrayMap<>();
-    private ArrayMap<IBinder, Integer> mLocalColorCallbackDisplayId = new ArrayMap<>();
+    private LocalColorRepository mLocalColorRepo = new LocalColorRepository();
 
     static class WallpaperData {
 
@@ -1305,24 +1298,16 @@ public class WallpaperManagerService extends IWallpaperManager.Stub
         public void onLocalWallpaperColorsChanged(RectF area, WallpaperColors colors,
                 int displayId) {
             forEachDisplayConnector(displayConnector -> {
-                if (displayConnector.mDisplayId == displayId) {
-                    RemoteCallbackList<ILocalWallpaperColorConsumer> callbacks;
-                    ArrayMap<IBinder, Integer> callbackDisplayIds;
-                    synchronized (mLock) {
-                        callbacks = mLocalColorAreaCallbacks.get(area);
-                        callbackDisplayIds = new ArrayMap<>(mLocalColorCallbackDisplayId);
+                Consumer<ILocalWallpaperColorConsumer> callback = cb -> {
+                    try {
+                        cb.onColorsChanged(area, colors);
+                    } catch (RemoteException e) {
+                        e.printStackTrace();
                     }
-                    if (callbacks == null) return;
-                    callbacks.broadcast(c -> {
-                        try {
-                            Integer targetDisplayId =
-                                    callbackDisplayIds.get(c.asBinder());
-                            if (targetDisplayId == null) return;
-                            if (targetDisplayId == displayId) c.onColorsChanged(area, colors);
-                        } catch (RemoteException e) {
-                            e.printStackTrace();
-                        }
-                    });
+                };
+                synchronized (mLock) {
+                    // it is safe to make an IPC call since it is one way (returns immediately)
+                    mLocalColorRepo.forEachCallback(callback, area, displayId);
                 }
             });
         }
@@ -1491,10 +1476,10 @@ public class WallpaperManagerService extends IWallpaperManager.Stub
                     Slog.w(TAG, "Failed to request wallpaper colors", e);
                 }
 
-                ArraySet<RectF> areas = mLocalColorDisplayIdAreas.get(displayId);
+                List<RectF> areas = mLocalColorRepo.getAreasByDisplayId(displayId);
                 if (areas != null && areas.size() != 0) {
                     try {
-                        connector.mEngine.addLocalColorsAreas(new ArrayList<>(areas));
+                        connector.mEngine.addLocalColorsAreas(areas);
                     } catch (RemoteException e) {
                         Slog.w(TAG, "Failed to register local colors areas", e);
                     }
@@ -2494,37 +2479,10 @@ public class WallpaperManagerService extends IWallpaperManager.Stub
         }
         IWallpaperEngine engine = getEngine(which, userId, displayId);
         if (engine == null) return;
-        ArrayList<RectF> validAreas = new ArrayList<>(regions.size());
         synchronized (mLock) {
-            ArraySet<RectF> areas = mLocalColorCallbackAreas.get(callback);
-            if (areas == null) areas = new ArraySet<>(regions.size());
-            areas.addAll(regions);
-            mLocalColorCallbackAreas.put(callback.asBinder(), areas);
+            mLocalColorRepo.addAreas(callback, regions, displayId);
         }
-        for (int i = 0; i < regions.size(); i++) {
-            if (!LOCAL_COLOR_BOUNDS.contains(regions.get(i))) {
-                continue;
-            }
-            RemoteCallbackList callbacks;
-            synchronized (mLock) {
-                callbacks = mLocalColorAreaCallbacks.get(
-                        regions.get(i));
-                if (callbacks == null) {
-                    callbacks = new RemoteCallbackList();
-                    mLocalColorAreaCallbacks.put(regions.get(i), callbacks);
-                }
-                mLocalColorCallbackDisplayId.put(callback.asBinder(), displayId);
-                ArraySet<RectF> displayAreas = mLocalColorDisplayIdAreas.get(displayId);
-                if (displayAreas == null) {
-                    displayAreas = new ArraySet<>(1);
-                    mLocalColorDisplayIdAreas.put(displayId, displayAreas);
-                }
-                displayAreas.add(regions.get(i));
-            }
-            validAreas.add(regions.get(i));
-            callbacks.register(callback);
-        }
-        engine.addLocalColorsAreas(validAreas);
+        engine.addLocalColorsAreas(regions);
     }
 
     @Override
@@ -2539,45 +2497,19 @@ public class WallpaperManagerService extends IWallpaperManager.Stub
             throw new SecurityException("calling user id does not match");
         }
         final long identity = Binder.clearCallingIdentity();
-        ArrayList<RectF> purgeAreas = new ArrayList<>();
-        IBinder binder = callback.asBinder();
+        List<RectF> purgeAreas = null;
         try {
             synchronized (mLock) {
-                ArraySet<RectF> currentAreas = mLocalColorCallbackAreas.get(binder);
-                if (currentAreas == null) return;
-                currentAreas.removeAll(removeAreas);
-                if (currentAreas.size() == 0) {
-                    mLocalColorCallbackDisplayId.remove(binder);
-                    for (RectF removeArea : removeAreas) {
-                        RemoteCallbackList<ILocalWallpaperColorConsumer> remotes =
-                                mLocalColorAreaCallbacks.get(removeArea);
-                        if (remotes == null) continue;
-                        remotes.unregister(callback);
-                        if (remotes.getRegisteredCallbackCount() == 0) {
-                            mLocalColorAreaCallbacks.remove(removeArea);
-                            purgeAreas.add(removeArea);
-                            ArraySet<RectF> displayAreas = mLocalColorDisplayIdAreas.get(displayId);
-                            if (displayAreas != null) {
-                                displayAreas.remove(removeArea);
-                                if (displayAreas.size() == 0) {
-                                    mLocalColorDisplayIdAreas.remove(displayId);
-                                }
-                            }
-                        }
-                    }
-                }
+                purgeAreas = mLocalColorRepo.removeAreas(callback, removeAreas, displayId);
             }
-
         } catch (Exception e) {
             // ignore any exception
         } finally {
             Binder.restoreCallingIdentity(identity);
         }
-
-        if (purgeAreas.size() == 0) return;
         IWallpaperEngine engine = getEngine(which, userId, displayId);
-        if (engine == null) return;
-        engine.removeLocalColorsAreas(purgeAreas);
+        if (engine == null || purgeAreas == null) return;
+        if (purgeAreas.size() > 0) engine.removeLocalColorsAreas(purgeAreas);
     }
 
     @Override
diff --git a/services/tests/mockingservicestests/src/com/android/server/wallpaper/LocalColorRepositoryTest.java b/services/tests/mockingservicestests/src/com/android/server/wallpaper/LocalColorRepositoryTest.java
new file mode 100644
index 000000000000..ea7a9a45facd
--- /dev/null
+++ b/services/tests/mockingservicestests/src/com/android/server/wallpaper/LocalColorRepositoryTest.java
@@ -0,0 +1,135 @@
+/*
+ * Copyright (C) 2021 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.server.wallpaper;
+
+import static junit.framework.Assert.assertEquals;
+import static junit.framework.Assert.assertFalse;
+import static junit.framework.Assert.assertTrue;
+
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.when;
+import static org.mockito.MockitoAnnotations.initMocks;
+
+import static java.util.Arrays.asList;
+
+import android.app.ILocalWallpaperColorConsumer;
+import android.app.WallpaperColors;
+import android.graphics.RectF;
+import android.os.IBinder;
+
+import androidx.test.runner.AndroidJUnit4;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+
+import android.util.ArraySet;
+import java.util.List;
+import java.util.function.Consumer;
+
+
+@RunWith(AndroidJUnit4.class)
+public class LocalColorRepositoryTest {
+    private LocalColorRepository mRepo = new LocalColorRepository();
+    @Mock
+    private IBinder mBinder1;
+    @Mock
+    private IBinder mBinder2;
+    @Mock
+    private ILocalWallpaperColorConsumer mCallback1;
+    @Mock
+    private ILocalWallpaperColorConsumer mCallback2;
+
+    @Before
+    public void setUp() {
+        initMocks(this);
+        when(mCallback1.asBinder()).thenReturn(mBinder1);
+        when(mCallback2.asBinder()).thenReturn(mBinder2);
+    }
+
+    @Test
+    public void testDisplayAreas() {
+        RectF area1 = new RectF(1, 0, 0, 0);
+        RectF area2 = new RectF(2, 1, 1, 1);
+        ArraySet<RectF> expectedAreas = new ArraySet(asList(area1, area2));
+
+        mRepo.addAreas(mCallback1, asList(area1), 0);
+        mRepo.addAreas(mCallback2, asList(area2), 0);
+        mRepo.addAreas(mCallback1, asList(new RectF(3, 1, 1, 1)), 1);
+
+        assertEquals(expectedAreas, new ArraySet(mRepo.getAreasByDisplayId(0)));
+        assertEquals(new ArraySet(asList(new RectF(3, 1, 1, 1))),
+                new ArraySet(mRepo.getAreasByDisplayId(1)));
+        assertEquals(new ArraySet(), new ArraySet(mRepo.getAreasByDisplayId(2)));
+    }
+
+    @Test
+    public void testAddAndRemoveAreas() {
+        RectF area1 = new RectF(1, 0, 0, 0);
+        RectF area2 = new RectF(2, 1, 1, 1);
+
+        mRepo.addAreas(mCallback1, asList(area1), 0);
+        mRepo.addAreas(mCallback1, asList(area2), 0);
+        mRepo.addAreas(mCallback2, asList(area2), 1);
+
+        List<RectF> removed = mRepo.removeAreas(mCallback1, asList(area1), 0);
+        assertEquals(new ArraySet(asList(area1)), new ArraySet(removed));
+        // since we have another callback with a different area, we don't purge rid of any areas
+        removed = mRepo.removeAreas(mCallback1, asList(area2), 0);
+        assertEquals(new ArraySet(), new ArraySet(removed));
+    }
+
+    @Test
+    public void testAreaCallback() {
+        Consumer<ILocalWallpaperColorConsumer> consumer = mock(Consumer.class);
+        WallpaperColors colors = mock(WallpaperColors.class);
+        RectF area1 = new RectF(1, 0, 0, 0);
+        RectF area2 = new RectF(2, 1, 1, 1);
+
+        mRepo.addAreas(mCallback1, asList(area1), 0);
+        mRepo.addAreas(mCallback1, asList(area2), 0);
+        mRepo.addAreas(mCallback2, asList(area2), 0);
+
+        mRepo.forEachCallback(consumer, area1, 0);
+        Mockito.verify(consumer, times(1)).accept(eq(mCallback1));
+        Mockito.verify(consumer, times(0)).accept(eq(mCallback2));
+        mRepo.forEachCallback(consumer, area2, 0);
+        Mockito.verify(consumer, times(2)).accept(eq(mCallback1));
+        Mockito.verify(consumer, times(1)).accept(eq(mCallback2));
+    }
+
+    @Test
+    public void unregisterCallbackWhenNoAreas() {
+        RectF area1 = new RectF(1, 0, 0, 0);
+        RectF area2 = new RectF(2, 1, 1, 1);
+
+        assertFalse(mRepo.isCallbackAvailable(mCallback1));
+
+        mRepo.addAreas(mCallback1, asList(area1), 0);
+        mRepo.addAreas(mCallback1, asList(area2), 0);
+
+        mRepo.removeAreas(mCallback1, asList(area1, area2), 0);
+        assertFalse(mRepo.isCallbackAvailable(mCallback1));
+
+        mRepo.addAreas(mCallback1, asList(area1), 0);
+        assertTrue(mRepo.isCallbackAvailable(mCallback1));
+    }
+}