Use a BgExecutor in SecurityControllerImpl

For more robust tests! Now, we can make sure the executor runs all
its runnables instead of waiting (test was previously very flaky).

Test: atest SystemUITests
Fixes: 149204632
Change-Id: I8ad047538b31709837c68399dcc1ee5b950283f9
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
index 312c4ac..d29f4fc 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/policy/SecurityControllerImpl.java
@@ -31,14 +31,12 @@
 import android.net.Network;
 import android.net.NetworkCapabilities;
 import android.net.NetworkRequest;
-import android.os.AsyncTask;
 import android.os.Handler;
 import android.os.RemoteException;
 import android.os.ServiceManager;
 import android.os.UserHandle;
 import android.os.UserManager;
 import android.security.KeyChain;
-import android.security.KeyChain.KeyChainConnection;
 import android.util.ArrayMap;
 import android.util.Log;
 import android.util.Pair;
@@ -55,6 +53,7 @@
 import java.io.FileDescriptor;
 import java.io.PrintWriter;
 import java.util.ArrayList;
+import java.util.concurrent.Executor;
 
 import javax.inject.Inject;
 import javax.inject.Singleton;
@@ -85,7 +84,7 @@
     private final DevicePolicyManager mDevicePolicyManager;
     private final PackageManager mPackageManager;
     private final UserManager mUserManager;
-    private final Handler mBgHandler;
+    private final Executor mBgExecutor;
 
     @GuardedBy("mCallbacks")
     private final ArrayList<SecurityControllerCallback> mCallbacks = new ArrayList<>();
@@ -101,16 +100,14 @@
     /**
      */
     @Inject
-    public SecurityControllerImpl(Context context, @Background Handler bgHandler,
-            BroadcastDispatcher broadcastDispatcher) {
-        this(context, bgHandler, broadcastDispatcher, null);
-    }
-
-    public SecurityControllerImpl(Context context, Handler bgHandler,
-            BroadcastDispatcher broadcastDispatcher, SecurityControllerCallback callback) {
+    public SecurityControllerImpl(
+            Context context,
+            @Background Handler bgHandler,
+            BroadcastDispatcher broadcastDispatcher,
+            @Background Executor bgExecutor
+    ) {
         super(broadcastDispatcher);
         mContext = context;
-        mBgHandler = bgHandler;
         mDevicePolicyManager = (DevicePolicyManager)
                 context.getSystemService(Context.DEVICE_POLICY_SERVICE);
         mConnectivityManager = (ConnectivityManager)
@@ -118,10 +115,8 @@
         mConnectivityManagerService = IConnectivityManager.Stub.asInterface(
                 ServiceManager.getService(Context.CONNECTIVITY_SERVICE));
         mPackageManager = context.getPackageManager();
-        mUserManager = (UserManager)
-                context.getSystemService(Context.USER_SERVICE);
-
-        addCallback(callback);
+        mUserManager = (UserManager) context.getSystemService(Context.USER_SERVICE);
+        mBgExecutor = bgExecutor;
 
         IntentFilter filter = new IntentFilter();
         filter.addAction(KeyChain.ACTION_TRUST_STORE_CHANGED);
@@ -305,7 +300,23 @@
     }
 
     private void refreshCACerts(int userId) {
-        new CACertLoader().execute(userId);
+        mBgExecutor.execute(() -> {
+            Pair<Integer, Boolean> idWithCert = null;
+            try (KeyChain.KeyChainConnection conn = KeyChain.bindAsUser(mContext,
+                    UserHandle.of(userId))) {
+                boolean hasCACerts = !(conn.getService().getUserCaAliases().getList().isEmpty());
+                idWithCert = new Pair<Integer, Boolean>(userId, hasCACerts);
+            } catch (RemoteException | InterruptedException | AssertionError e) {
+                Log.i(TAG, "failed to get CA certs", e);
+                idWithCert = new Pair<Integer, Boolean>(userId, null);
+            } finally {
+                if (DEBUG) Log.d(TAG, "Refreshing CA Certs " + idWithCert);
+                if (idWithCert != null && idWithCert.second != null) {
+                    mHasCACerts.put(idWithCert.first, idWithCert.second);
+                    fireCallbacks();
+                }
+            }
+        });
     }
 
     private String getNameForVpnConfig(VpnConfig cfg, UserHandle user) {
@@ -408,28 +419,4 @@
             }
         }
     };
-
-    protected class CACertLoader extends AsyncTask<Integer, Void, Pair<Integer, Boolean> > {
-
-        @Override
-        protected Pair<Integer, Boolean> doInBackground(Integer... userId) {
-            try (KeyChainConnection conn = KeyChain.bindAsUser(mContext,
-                                                               UserHandle.of(userId[0]))) {
-                boolean hasCACerts = !(conn.getService().getUserCaAliases().getList().isEmpty());
-                return new Pair<Integer, Boolean>(userId[0], hasCACerts);
-            } catch (RemoteException | InterruptedException | AssertionError e) {
-                Log.i(TAG, "failed to get CA certs", e);
-                return new Pair<Integer, Boolean>(userId[0], null);
-            }
-        }
-
-        @Override
-        protected void onPostExecute(Pair<Integer, Boolean> result) {
-            if (DEBUG) Log.d(TAG, "onPostExecute " + result);
-            if (result.second != null) {
-                mHasCACerts.put(result.first, result.second);
-                fireCallbacks();
-            }
-        }
-    }
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
index e6b0440..44184ee 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/policy/SecurityControllerTest.java
@@ -21,6 +21,7 @@
 import static org.junit.Assert.assertTrue;
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyInt;
+import static org.mockito.Matchers.anyObject;
 import static org.mockito.Matchers.argThat;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
@@ -28,6 +29,7 @@
 import static org.mockito.Mockito.when;
 
 import android.app.admin.DevicePolicyManager;
+import android.content.BroadcastReceiver;
 import android.content.ComponentName;
 import android.content.Context;
 import android.content.Intent;
@@ -37,7 +39,6 @@
 import android.net.ConnectivityManager.NetworkCallback;
 import android.net.NetworkRequest;
 import android.os.Handler;
-import android.os.Looper;
 import android.os.UserManager;
 import android.security.IKeyChainService;
 import android.test.suitebuilder.annotation.SmallTest;
@@ -46,34 +47,30 @@
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.broadcast.BroadcastDispatcher;
-import com.android.systemui.statusbar.policy.SecurityController.SecurityControllerCallback;
+import com.android.systemui.util.concurrency.FakeExecutor;
+import com.android.systemui.util.time.FakeSystemClock;
 
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
-import java.util.concurrent.CountDownLatch;
-import java.util.concurrent.TimeUnit;
 
 @SmallTest
 @RunWith(AndroidJUnit4.class)
-public class SecurityControllerTest extends SysuiTestCase implements SecurityControllerCallback {
+public class SecurityControllerTest extends SysuiTestCase {
     private final DevicePolicyManager mDevicePolicyManager = mock(DevicePolicyManager.class);
     private final IKeyChainService.Stub mKeyChainService = mock(IKeyChainService.Stub.class);
     private final UserManager mUserManager = mock(UserManager.class);
+    private final BroadcastDispatcher mBroadcastDispatcher = mock(BroadcastDispatcher.class);
+    private final Handler mHandler = mock(Handler.class);
     private SecurityControllerImpl mSecurityController;
-    private CountDownLatch mStateChangedLatch;
     private ConnectivityManager mConnectivityManager = mock(ConnectivityManager.class);
-
-    // implementing SecurityControllerCallback
-    @Override
-    public void onStateChanged() {
-        mStateChangedLatch.countDown();
-    }
+    private FakeExecutor mBgExecutor;
+    private BroadcastReceiver mBroadcastReceiver;
 
     @Before
     public void setUp() throws Exception {
@@ -95,18 +92,23 @@
         when(mKeyChainService.queryLocalInterface("android.security.IKeyChainService"))
                 .thenReturn(mKeyChainService);
 
-        // Wait for callbacks from the onUserSwitched() function in the
-        // constructor of mSecurityController
-        mStateChangedLatch = new CountDownLatch(1);
-        // TODO: Migrate this test to TestableLooper and use a handler attached
-        // to that.
-        mSecurityController = new SecurityControllerImpl(mContext,
-                new Handler(Looper.getMainLooper()), mock(BroadcastDispatcher.class), this);
-    }
+        ArgumentCaptor<BroadcastReceiver> brCaptor =
+                ArgumentCaptor.forClass(BroadcastReceiver.class);
 
-    @After
-    public void tearDown() {
-        mSecurityController.removeCallback(this);
+        mBgExecutor = new FakeExecutor(new FakeSystemClock());
+        mSecurityController = new SecurityControllerImpl(
+                mContext,
+                mHandler,
+                mBroadcastDispatcher,
+                mBgExecutor);
+
+        verify(mBroadcastDispatcher).registerReceiverWithHandler(
+                brCaptor.capture(),
+                anyObject(),
+                anyObject(),
+                anyObject());
+
+        mBroadcastReceiver = brCaptor.getValue();
     }
 
     @Test
@@ -126,8 +128,6 @@
 
     @Test
     public void testWorkAccount() throws Exception {
-        // Wait for the callbacks from setUp()
-        assertTrue(mStateChangedLatch.await(1, TimeUnit.SECONDS));
         assertFalse(mSecurityController.hasCACertInCurrentUser());
 
         final int PRIMARY_USER_ID = 0;
@@ -140,53 +140,41 @@
         assertTrue(mSecurityController.hasWorkProfile());
         assertFalse(mSecurityController.hasCACertInWorkProfile());
 
-        mStateChangedLatch = new CountDownLatch(1);
-
         when(mKeyChainService.getUserCaAliases())
                 .thenReturn(new StringParceledListSlice(Arrays.asList("One CA Alias")));
 
-        mSecurityController.new CACertLoader()
-                           .execute(MANAGED_USER_ID);
+        refreshCACerts(MANAGED_USER_ID);
+        mBgExecutor.runAllReady();
 
-        assertTrue(mStateChangedLatch.await(3, TimeUnit.SECONDS));
         assertTrue(mSecurityController.hasCACertInWorkProfile());
     }
 
     @Test
     public void testCaCertLoader() throws Exception {
-        // Wait for the callbacks from setUp()
-        assertTrue(mStateChangedLatch.await(1, TimeUnit.SECONDS));
         assertFalse(mSecurityController.hasCACertInCurrentUser());
 
         // With a CA cert
-        mStateChangedLatch = new CountDownLatch(1);
-
         when(mKeyChainService.getUserCaAliases())
                 .thenReturn(new StringParceledListSlice(Arrays.asList("One CA Alias")));
 
-        mSecurityController.new CACertLoader()
-                           .execute(0);
+        refreshCACerts(0);
+        mBgExecutor.runAllReady();
 
-        assertTrue(mStateChangedLatch.await(3, TimeUnit.SECONDS));
         assertTrue(mSecurityController.hasCACertInCurrentUser());
 
         // Exception
-        mStateChangedLatch = new CountDownLatch(1);
-
         when(mKeyChainService.getUserCaAliases())
                 .thenThrow(new AssertionError("Test AssertionError"))
                 .thenReturn(new StringParceledListSlice(new ArrayList<String>()));
 
-        mSecurityController.new CACertLoader()
-                           .execute(0);
+        refreshCACerts(0);
+        mBgExecutor.runAllReady();
 
-        assertFalse(mStateChangedLatch.await(1, TimeUnit.SECONDS));
         assertTrue(mSecurityController.hasCACertInCurrentUser());
 
-        mSecurityController.new CACertLoader()
-                           .execute(0);
+        refreshCACerts(0);
+        mBgExecutor.runAllReady();
 
-        assertTrue(mStateChangedLatch.await(1, TimeUnit.SECONDS));
         assertFalse(mSecurityController.hasCACertInCurrentUser());
     }
 
@@ -197,4 +185,13 @@
                         && request.networkCapabilities.getCapabilities().length == 0
                 ), any(NetworkCallback.class));
     }
+
+    /**
+     * refresh CA certs by sending a user unlocked broadcast for the desired user
+     */
+    private void refreshCACerts(int userId) {
+        Intent intent = new Intent(Intent.ACTION_USER_UNLOCKED);
+        intent.putExtra(Intent.EXTRA_USER_HANDLE, userId);
+        mBroadcastReceiver.onReceive(mContext, intent);
+    }
 }