From b5b32d479d44c55922e06395861c21b0ea6bd024 Mon Sep 17 00:00:00 2001
From: Riddle Hsu <riddlehsu@google.com>
Date: Wed, 2 Mar 2022 22:14:20 -0600
Subject: [PATCH] Fix NPE in WakeLock#finalize for unit test

The stubbed method of "mock(PowerManager.WakeLock.class)" will
lose effect after clearInlineMocks. And then finalizer runs
the real method finalize() that cause NPE because the mocked
object didn't initialize fields.

Bug: 219640050
Test: atest WmTests

Change-Id: If8ee4fbc98c9e33d6ab6ba75a65ba177fa06f09c
---
 .../wm/ActivityTaskManagerServiceTests.java   |  3 ++-
 .../server/wm/SystemServicesTestRule.java     | 19 +++++++++++++++++--
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java b/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
index 741f33f8aca7e..87c75bef85aab 100644
--- a/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
+++ b/services/tests/wmtests/src/com/android/server/wm/ActivityTaskManagerServiceTests.java
@@ -267,7 +267,8 @@ public class ActivityTaskManagerServiceTests extends WindowTestsBase {
     @Test
     public void testUpdateSleep() {
         doCallRealMethod().when(mWm.mRoot).hasAwakeDisplay();
-        mSupervisor.mGoingToSleepWakeLock = mock(PowerManager.WakeLock.class);
+        mSupervisor.mGoingToSleepWakeLock =
+                mSystemServicesTestRule.createStubbedWakeLock(true /* needVerification */);
         final Task rootHomeTask = mWm.mRoot.getDefaultTaskDisplayArea().getOrCreateRootHomeTask();
         final ActivityRecord homeActivity = new ActivityBuilder(mAtm).setTask(rootHomeTask).build();
         final ActivityRecord topActivity = new ActivityBuilder(mAtm).setCreateTask(true).build();
diff --git a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
index 8e7ba4bc3293f..0b2399b18239d 100644
--- a/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
+++ b/services/tests/wmtests/src/com/android/server/wm/SystemServicesTestRule.java
@@ -102,6 +102,10 @@ public class SystemServicesTestRule implements TestRule {
     static int sNextTaskId = 100;
 
     private static final int[] TEST_USER_PROFILE_IDS = {};
+    /** Use a real static object so there won't be NPE in finalize() after clearInlineMocks(). */
+    private static final PowerManager.WakeLock sWakeLock = getInstrumentation().getContext()
+            .getSystemService(PowerManager.class).newWakeLock(PowerManager.PARTIAL_WAKE_LOCK, TAG);
+    private PowerManager.WakeLock mStubbedWakeLock;
 
     private Description mDescription;
     private Context mContext;
@@ -196,7 +200,8 @@ public class SystemServicesTestRule implements TestRule {
         // Prevent "WakeLock finalized while still held: SCREEN_FROZEN".
         final PowerManager pm = mock(PowerManager.class);
         doReturn(pm).when(mContext).getSystemService(eq(Context.POWER_SERVICE));
-        doReturn(mock(PowerManager.WakeLock.class)).when(pm).newWakeLock(anyInt(), anyString());
+        mStubbedWakeLock = createStubbedWakeLock(false /* needVerification */);
+        doReturn(mStubbedWakeLock).when(pm).newWakeLock(anyInt(), anyString());
 
         // DisplayManagerInternal
         final DisplayManagerInternal dmi = mock(DisplayManagerInternal.class);
@@ -403,6 +408,16 @@ public class SystemServicesTestRule implements TestRule {
         return mPowerManagerWrapper;
     }
 
+    /** Creates a no-op wakelock object. */
+    PowerManager.WakeLock createStubbedWakeLock(boolean needVerification) {
+        if (needVerification) {
+            return mock(PowerManager.WakeLock.class, Mockito.withSettings()
+                    .spiedInstance(sWakeLock).defaultAnswer(Mockito.RETURNS_DEFAULTS));
+        }
+        return mock(PowerManager.WakeLock.class, Mockito.withSettings()
+                .spiedInstance(sWakeLock).stubOnly());
+    }
+
     void setSurfaceFactory(Supplier<Surface> factory) {
         mSurfaceFactory = factory;
     }
@@ -558,7 +573,7 @@ public class SystemServicesTestRule implements TestRule {
             // unit test version does not handle launch wake lock
             doNothing().when(this).acquireLaunchWakelock();
 
-            mLaunchingActivityWakeLock = mock(PowerManager.WakeLock.class);
+            mLaunchingActivityWakeLock = mStubbedWakeLock;
 
             initialize();