From 3536b70f119b58d636e5e2fa0154cdf3692a9759 Mon Sep 17 00:00:00 2001
From: Bryce Lee <brycelee@google.com>
Date: Wed, 16 Mar 2022 12:56:55 -0700
Subject: [PATCH] Remove complication id from entry mapping on removal.

This changelist ensures a complication's id is removed from the
ComplicationLayoutEngine's internal entry mapping when the complication
is removed. This can happen when complications are singletons.

Bug: 224809882
Test: atest ComplicationLayoutEngineTest#testDoubleRemoval
Change-Id: Ie350e7ab08e909af43ad9d7eadd0ca82c6ad3aa0
---
 .../ComplicationLayoutEngine.java             | 10 +++---
 .../ComplicationLayoutEngineTest.java         | 34 +++++++++++++++++++
 2 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/packages/SystemUI/src/com/android/systemui/dreams/complication/ComplicationLayoutEngine.java b/packages/SystemUI/src/com/android/systemui/dreams/complication/ComplicationLayoutEngine.java
index aa433832d4623..4065a25dac803 100644
--- a/packages/SystemUI/src/com/android/systemui/dreams/complication/ComplicationLayoutEngine.java
+++ b/packages/SystemUI/src/com/android/systemui/dreams/complication/ComplicationLayoutEngine.java
@@ -540,13 +540,15 @@ public class ComplicationLayoutEngine implements Complication.VisibilityControll
     /**
      * Removes a complication by {@link ComplicationId}.
      */
-    public void removeComplication(ComplicationId id) {
-        if (!mEntries.containsKey(id)) {
+    public boolean removeComplication(ComplicationId id) {
+        final ViewEntry entry = mEntries.remove(id);
+
+        if (entry == null) {
             Log.e(TAG, "could not find id:" + id);
-            return;
+            return false;
         }
 
-        final ViewEntry entry = mEntries.get(id);
         entry.remove();
+        return true;
     }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/dreams/complication/ComplicationLayoutEngineTest.java b/packages/SystemUI/tests/src/com/android/systemui/dreams/complication/ComplicationLayoutEngineTest.java
index d1d9ec3af711f..35bcfcd39fca5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/dreams/complication/ComplicationLayoutEngineTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/dreams/complication/ComplicationLayoutEngineTest.java
@@ -18,6 +18,7 @@ package com.android.systemui.dreams.complication;
 import static com.google.common.truth.Truth.assertThat;
 
 import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
@@ -414,4 +415,37 @@ public class ComplicationLayoutEngineTest extends SysuiTestCase {
             assertThat(lp.endToEnd == ConstraintLayout.LayoutParams.PARENT_ID).isTrue();
         });
     }
+
+    /**
+     * Ensures a second removal of a complication is a no-op.
+     */
+    @Test
+    public void testDoubleRemoval() {
+        final ComplicationLayoutEngine engine =
+                new ComplicationLayoutEngine(mLayout, 0, mTouchSession, 0, 0);
+
+        final ViewInfo firstViewInfo = new ViewInfo(
+                new ComplicationLayoutParams(
+                        100,
+                        100,
+                        ComplicationLayoutParams.POSITION_TOP
+                                | ComplicationLayoutParams.POSITION_END,
+                        ComplicationLayoutParams.DIRECTION_DOWN,
+                        0),
+                Complication.CATEGORY_STANDARD,
+                mLayout);
+
+        engine.addComplication(firstViewInfo.id, firstViewInfo.view, firstViewInfo.lp,
+                firstViewInfo.category);
+        verify(mLayout).addView(firstViewInfo.view);
+
+        assertThat(engine.removeComplication(firstViewInfo.id)).isTrue();
+        verify(firstViewInfo.view).getParent();
+        verify(mLayout).removeView(firstViewInfo.view);
+
+        Mockito.clearInvocations(mLayout, firstViewInfo.view);
+        assertThat(engine.removeComplication(firstViewInfo.id)).isFalse();
+        verify(firstViewInfo.view, never()).getParent();
+        verify(mLayout, never()).removeView(firstViewInfo.view);
+    }
 }