diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/leak/LeakDetectorTest.java b/packages/SystemUI/tests/src/com/android/systemui/util/leak/LeakDetectorTest.java
index f1965a223dbc7..56de32d3a4010 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/leak/LeakDetectorTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/leak/LeakDetectorTest.java
@@ -33,6 +33,8 @@ import java.io.FileOutputStream;
 import java.io.PrintWriter;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
@@ -40,6 +42,30 @@ public class LeakDetectorTest extends SysuiTestCase {
 
     private LeakDetector mLeakDetector;
 
+    // The references for which collection is observed are stored in fields. The allocation and
+    // of these references happens in separate methods (trackObjectWith/trackCollectionWith)
+    // from where they are set to null. The generated code might keep the allocated reference
+    // alive in a dex register when compiling in release mode. As R8 is used to compile this
+    // test the --dontoptimize flag is also required to ensure that these methods are not
+    // inlined, as that would defeat the purpose of having the mutation in methods.
+    private Object mObject;
+    private Collection<?> mCollection;
+
+    private CollectionWaiter trackObjectWith(Consumer<Object> tracker) {
+        mObject = new Object();
+        CollectionWaiter collectionWaiter = ReferenceTestUtils.createCollectionWaiter(mObject);
+        tracker.accept(mObject);
+        return collectionWaiter;
+    }
+
+    private CollectionWaiter trackCollectionWith(
+            BiConsumer<? super Collection<?>, String> tracker) {
+        mCollection = new ArrayList<>();
+        CollectionWaiter collectionWaiter = ReferenceTestUtils.createCollectionWaiter(mCollection);
+        tracker.accept(mCollection, "tag");
+        return collectionWaiter;
+    }
+
     @Before
     public void setup() {
         mLeakDetector = LeakDetector.create();
@@ -51,31 +77,22 @@ public class LeakDetectorTest extends SysuiTestCase {
 
     @Test
     public void trackInstance_doesNotLeakTrackedObject() {
-        Object object = new Object();
-        CollectionWaiter collectionWaiter = ReferenceTestUtils.createCollectionWaiter(object);
-
-        mLeakDetector.trackInstance(object);
-        object = null;
+        CollectionWaiter collectionWaiter = trackObjectWith(mLeakDetector::trackInstance);
+        mObject = null;
         collectionWaiter.waitForCollection();
     }
 
     @Test
     public void trackCollection_doesNotLeakTrackedObject() {
-        Collection<?> object = new ArrayList<>();
-        CollectionWaiter collectionWaiter = ReferenceTestUtils.createCollectionWaiter(object);
-
-        mLeakDetector.trackCollection(object, "tag");
-        object = null;
+        CollectionWaiter collectionWaiter = trackCollectionWith(mLeakDetector::trackCollection);
+        mCollection = null;
         collectionWaiter.waitForCollection();
     }
 
     @Test
     public void trackGarbage_doesNotLeakTrackedObject() {
-        Object object = new Object();
-        CollectionWaiter collectionWaiter = ReferenceTestUtils.createCollectionWaiter(object);
-
-        mLeakDetector.trackGarbage(object);
-        object = null;
+        CollectionWaiter collectionWaiter = trackObjectWith(mLeakDetector::trackGarbage);
+        mObject = null;
         collectionWaiter.waitForCollection();
     }
 
@@ -108,4 +125,4 @@ public class LeakDetectorTest extends SysuiTestCase {
         FileOutputStream fos = new FileOutputStream("/dev/null");
         mLeakDetector.dump(fos.getFD(), new PrintWriter(fos), new String[0]);
     }
-}
\ No newline at end of file
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/util/leak/WeakIdentityHashMapTest.java b/packages/SystemUI/tests/src/com/android/systemui/util/leak/WeakIdentityHashMapTest.java
index 9787df91f5dee..ce6212ef5aae5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/util/leak/WeakIdentityHashMapTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/util/leak/WeakIdentityHashMapTest.java
@@ -41,6 +41,13 @@ public class WeakIdentityHashMapTest extends SysuiTestCase {
         mMap = new WeakIdentityHashMap<>();
     }
 
+    private CollectionWaiter addObjectToMap(WeakIdentityHashMap<Object, Object> map) {
+      Object object = new Object();
+      CollectionWaiter collectionWaiter = ReferenceTestUtils.createCollectionWaiter(object);
+      map.put(object, "value");
+      return collectionWaiter;
+    }
+
     @Test
     public void testUsesIdentity() {
         String a1 = new String("a");
@@ -56,11 +63,12 @@ public class WeakIdentityHashMapTest extends SysuiTestCase {
 
     @Test
     public void testWeaklyReferences() {
-        Object object = new Object();
-        CollectionWaiter collectionWaiter = ReferenceTestUtils.createCollectionWaiter(object);
-
-        mMap.put(object, "value");
-        object = null;
+        // Allocate and add an object to the weak map in a separate method to avoid a live
+        // reference to the allocated object in a dex register. As R8 is used to compile this
+        // test the --dontoptimize flag is also required to ensure that the method is not
+        // inlined, as that would defeat the purpose of having the allocation in a separate
+        // method.
+        CollectionWaiter collectionWaiter = addObjectToMap(mMap);
 
         // Wait until object has been collected. We'll also need to wait for mMap to become empty,
         // because our collection waiter may be told about the collection earlier than mMap.
@@ -70,4 +78,4 @@ public class WeakIdentityHashMapTest extends SysuiTestCase {
         assertEquals(0, mMap.size());
         assertTrue(mMap.isEmpty());
     }
-}
\ No newline at end of file
+}