From 7ff82b0a92b0a3115f16850352415ffbd34a3faa Mon Sep 17 00:00:00 2001
From: Lucas Dupin <dupin@google.com>
Date: Wed, 16 May 2018 12:14:10 -0700
Subject: [PATCH] Fix KeyguardUpdateMonitor test

KeyguardUpdateMonitor's Handler runs on the main looper by design,
we need to ensure that whenever we're processing messages, that
they will also be received on the main looper.

Also unregistered a broadcast listener to avoid a possible race
condition

Test: atest packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
Bug: 79550837
Change-Id: I14a319da5c6bc46fd32675ae205e14a6228efaa4
---
 .../keyguard/KeyguardUpdateMonitor.java       |  2 +-
 .../keyguard/KeyguardUpdateMonitorTest.java   | 47 ++++++++++++++-----
 2 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
index 9f382b00d4f87..31fd47ff8f47b 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
@@ -750,7 +750,7 @@ public class KeyguardUpdateMonitor implements TrustManager.TrustListener {
     private DisplayClientState mDisplayClientState = new DisplayClientState();
 
     @VisibleForTesting
-    final BroadcastReceiver mBroadcastReceiver = new BroadcastReceiver() {
+    protected final BroadcastReceiver mBroadcastReceiver = new BroadcastReceiver() {
 
         @Override
         public void onReceive(Context context, Intent intent) {
diff --git a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
index 21483aac2ff21..d46a9747ac240 100644
--- a/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
+++ b/packages/SystemUI/tests/src/com/android/keyguard/KeyguardUpdateMonitorTest.java
@@ -16,6 +16,7 @@
 
 package com.android.keyguard;
 
+import android.content.Context;
 import android.content.Intent;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.testing.AndroidTestingRunner;
@@ -35,7 +36,12 @@ import java.util.concurrent.atomic.AtomicBoolean;
 
 @SmallTest
 @RunWith(AndroidTestingRunner.class)
-@RunWithLooper
+// We must run on the main looper because KeyguardUpdateMonitor#mHandler is initialized with
+// new Handler(Looper.getMainLooper()).
+//
+// Using the main looper should be avoided whenever possible, please don't copy this over to
+// new tests.
+@RunWithLooper(setAsMainLooper = true)
 public class KeyguardUpdateMonitorTest extends SysuiTestCase {
 
     private TestableLooper mTestableLooper;
@@ -49,24 +55,39 @@ public class KeyguardUpdateMonitorTest extends SysuiTestCase {
     public void testIgnoresSimStateCallback_rebroadcast() {
         Intent intent = new Intent(TelephonyIntents.ACTION_SIM_STATE_CHANGED);
 
-        AtomicBoolean simStateChanged = new AtomicBoolean(false);
-        KeyguardUpdateMonitor keyguardUpdateMonitor = new KeyguardUpdateMonitor(getContext()) {
-            @Override
-            protected void handleSimStateChange(int subId, int slotId,
-                    IccCardConstants.State state) {
-                simStateChanged.set(true);
-                super.handleSimStateChange(subId, slotId, state);
-            }
-        };
+        TestableKeyguardUpdateMonitor keyguardUpdateMonitor =
+                new TestableKeyguardUpdateMonitor(getContext());
 
         keyguardUpdateMonitor.mBroadcastReceiver.onReceive(getContext(), intent);
         mTestableLooper.processAllMessages();
-        Assert.assertTrue("onSimStateChanged not called", simStateChanged.get());
+        Assert.assertTrue("onSimStateChanged not called",
+                keyguardUpdateMonitor.hasSimStateJustChanged());
 
         intent.putExtra(TelephonyIntents.EXTRA_REBROADCAST_ON_UNLOCK, true);
-        simStateChanged.set(false);
         keyguardUpdateMonitor.mBroadcastReceiver.onReceive(getContext(), intent);
         mTestableLooper.processAllMessages();
-        Assert.assertFalse("onSimStateChanged should have been skipped", simStateChanged.get());
+        Assert.assertFalse("onSimStateChanged should have been skipped",
+                keyguardUpdateMonitor.hasSimStateJustChanged());
+    }
+
+    private class TestableKeyguardUpdateMonitor extends KeyguardUpdateMonitor {
+        AtomicBoolean mSimStateChanged = new AtomicBoolean(false);
+
+        protected TestableKeyguardUpdateMonitor(Context context) {
+            super(context);
+            // Avoid race condition when unexpected broadcast could be received.
+            context.unregisterReceiver(mBroadcastReceiver);
+        }
+
+        public boolean hasSimStateJustChanged() {
+            return mSimStateChanged.getAndSet(false);
+        }
+
+        @Override
+        protected void handleSimStateChange(int subId, int slotId,
+                IccCardConstants.State state) {
+            mSimStateChanged.set(true);
+            super.handleSimStateChange(subId, slotId, state);
+        }
     }
 }