Merge "Reduce Repetitive Calls to TelephonyManager From Telephony2gUpdater"
diff --git a/src/com/android/phone/PhoneInterfaceManager.java b/src/com/android/phone/PhoneInterfaceManager.java
index bdbd5f9..ba8b330 100644
--- a/src/com/android/phone/PhoneInterfaceManager.java
+++ b/src/com/android/phone/PhoneInterfaceManager.java
@@ -2347,8 +2347,7 @@
         mRadioInterfaceCapabilities = RadioInterfaceCapabilityController.getInstance();
         mNotifyUserActivity = new AtomicBoolean(false);
         PropertyInvalidatedCache.invalidateCache(TelephonyManager.CACHE_KEY_PHONE_ACCOUNT_TO_SUBID);
-        mTelephony2gUpdater = new Telephony2gUpdater(
-                Executors.newSingleThreadExecutor(), mApp);
+        mTelephony2gUpdater = new Telephony2gUpdater(mApp);
         mTelephony2gUpdater.init();
         publish();
     }
diff --git a/src/com/android/phone/Telephony2gUpdater.java b/src/com/android/phone/Telephony2gUpdater.java
index 0919385..baaa684 100644
--- a/src/com/android/phone/Telephony2gUpdater.java
+++ b/src/com/android/phone/Telephony2gUpdater.java
@@ -30,8 +30,13 @@
 import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.telephony.RILConstants;
 
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
 
 /**
  * A {@link BroadcastReceiver} that ensures that user restrictions are correctly applied to
@@ -39,27 +44,50 @@
  * This includes handling broadcasts from user restriction state changes, as well as ensuring that
  * SIM-specific settings are correctly applied when new subscriptions become active.
  *
+ * <p>
  * Callers are expected to call {@code init()} and keep an instance of this class alive.
+ * </p>
  */
 public class Telephony2gUpdater extends BroadcastReceiver {
-    private static final String TAG = "TelephonyUserManagerReceiver";
+    private static final String TAG = "Telephony2gUpdater";
 
     // We can't interact with the HAL on the main thread of the phone process (where
     // receivers are run by default), so we execute our logic from a separate thread.
+    // The correctness of this implementation relies heavily on this executor ensuring
+    // tasks are serially executed i.e. ExecutorService.newSingleThreadExecutor()
     private final Executor mExecutor;
     private final Context mContext;
     private final long mBaseAllowedNetworks;
 
-    public Telephony2gUpdater(Executor executor, Context context) {
-        this(executor, context,
+    private UserManager mUserManager;
+    private TelephonyManager mTelephonyManager;
+    private SubscriptionManager mSubscriptionManager;
+
+    // The current subscription ids
+    // Ensure this value is never accessed concurrently
+    private Set<Integer> mCurrentSubscriptions;
+    // We keep track of the last value to avoid updating when unrelated user restrictions change
+    // Ensure this value is never accessed concurrently
+    private boolean mDisallowCellular2gRestriction;
+
+    public Telephony2gUpdater(Context context) {
+        this(Executors.newSingleThreadExecutor(), context,
                 RadioAccessFamily.getRafFromNetworkType(RILConstants.PREFERRED_NETWORK_MODE));
     }
 
-    public Telephony2gUpdater(Executor executor, Context context,
-            long baseAllowedNetworks) {
+    @VisibleForTesting
+    public Telephony2gUpdater(Executor executor, Context context, long baseAllowedNetworks) {
         mExecutor = executor;
         mContext = context;
         mBaseAllowedNetworks = baseAllowedNetworks;
+
+        mUserManager = mContext.getSystemService(UserManager.class);
+        mTelephonyManager = mContext.getSystemService(TelephonyManager.class);
+        mSubscriptionManager = mContext.getSystemService(SubscriptionManager.class);
+
+        // All user restrictions are false by default
+        mDisallowCellular2gRestriction = false;
+        mCurrentSubscriptions = new HashSet<>();
     }
 
     /**
@@ -80,41 +108,42 @@
         Log.i(TAG, "Received callback for action " + intent.getAction());
         final PendingResult result = goAsync();
         mExecutor.execute(() -> {
-            Log.i(TAG, "Running handler for action " + intent.getAction());
-            handleUserRestrictionsChanged(context);
-            result.finish();
+            boolean disallow2g = mUserManager.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G);
+            if (mDisallowCellular2gRestriction == disallow2g) {
+                Log.i(TAG, "No update to DISALLOW_CELLULAR_2G restriction.");
+                return;
+            }
+
+            mDisallowCellular2gRestriction = disallow2g;
+
+            Log.i(TAG, "Running handler for all subscriptions based on DISALLOW_CELLULAR_2G change."
+                    + " Restriction value: " + mDisallowCellular2gRestriction);
+            handleUserRestrictionsChanged(mCurrentSubscriptions);
+            if (result != null) {
+                result.finish();
+            }
         });
     }
 
     /**
-     * Update all active subscriptions with allowed network types depending on the current state
-     * of the {@link UserManager.DISALLOW_2G}.
+     * Update subscriptions with allowed network types depending on the current state
+     * of the {@link UserManager#DISALLOW_CELLULAR_2G}.
+     *
+     * @param subIds A list of subIds to update.
      */
-    @VisibleForTesting
-    public void handleUserRestrictionsChanged(Context context) {
-        UserManager um = context.getSystemService(UserManager.class);
-        TelephonyManager tm = context.getSystemService(TelephonyManager.class);
-        SubscriptionManager sm = context.getSystemService(SubscriptionManager.class);
+    private void handleUserRestrictionsChanged(Collection<Integer> subIds) {
         final long twoGBitmask = TelephonyManager.NETWORK_CLASS_BITMASK_2G;
 
-        boolean shouldDisable2g = um.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G);
-
-        // This is expected when subscription info cannot be determined. We'll get another
-        // callback in the future from our SubscriptionListener once we have valid subscriptions.
-        List<SubscriptionInfo> subscriptionInfoList = sm.getAvailableSubscriptionInfoList();
-        if (subscriptionInfoList == null) {
-            return;
-        }
-
         long allowedNetworkTypes = mBaseAllowedNetworks;
 
         // 2G device admin controls are global
-        for (SubscriptionInfo info : subscriptionInfoList) {
-            TelephonyManager telephonyManager = tm.createForSubscriptionId(
-                    info.getSubscriptionId());
-            if (shouldDisable2g) {
+        for (Integer subId : subIds) {
+            TelephonyManager telephonyManager = mTelephonyManager.createForSubscriptionId(subId);
+            if (mDisallowCellular2gRestriction) {
+                Log.i(TAG, "Disabling 2g based on user restriction for subId: " + subId);
                 allowedNetworkTypes &= ~twoGBitmask;
             } else {
+                Log.i(TAG, "Enabling 2g based on user restriction for subId: " + subId);
                 allowedNetworkTypes |= twoGBitmask;
             }
             telephonyManager.setAllowedNetworkTypesForReason(
@@ -126,8 +155,30 @@
     private class SubscriptionListener extends SubscriptionManager.OnSubscriptionsChangedListener {
         @Override
         public void onSubscriptionsChanged() {
-            Log.i(TAG, "Running handler for subscription change.");
-            handleUserRestrictionsChanged(mContext);
+            // Note that this entire callback gets invoked in the single threaded executor
+            List<SubscriptionInfo> allSubscriptions =
+                    mSubscriptionManager.getCompleteActiveSubscriptionInfoList();
+
+            HashSet<Integer> updatedSubIds = new HashSet<>(allSubscriptions.size());
+            List<Integer> newSubIds = new ArrayList<>();
+
+            for (SubscriptionInfo info : allSubscriptions) {
+                updatedSubIds.add(info.getSubscriptionId());
+                if (!mCurrentSubscriptions.contains(info.getSubscriptionId())) {
+                    newSubIds.add(info.getSubscriptionId());
+                }
+            }
+
+            mCurrentSubscriptions = updatedSubIds;
+
+            if (newSubIds.isEmpty()) {
+                Log.d(TAG, "No new subIds. Skipping update.");
+                return;
+            }
+
+            Log.i(TAG, "New subscriptions found. Running handler to update 2g restrictions with "
+                    + "subIds " + newSubIds.toString());
+            handleUserRestrictionsChanged(newSubIds);
         }
     }
 
diff --git a/tests/src/com/android/phone/Telephony2gUpdaterTest.java b/tests/src/com/android/phone/Telephony2gUpdaterTest.java
index 3443767..3684f30 100644
--- a/tests/src/com/android/phone/Telephony2gUpdaterTest.java
+++ b/tests/src/com/android/phone/Telephony2gUpdaterTest.java
@@ -16,13 +16,17 @@
 
 package com.android.phone;
 
+import static org.junit.Assert.fail;
+import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.anyInt;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
+import android.content.Intent;
 import android.os.UserManager;
 import android.telephony.SubscriptionInfo;
 import android.telephony.SubscriptionManager;
@@ -35,22 +39,33 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
 
 import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 @RunWith(AndroidJUnit4.class)
 public class Telephony2gUpdaterTest extends TelephonyTestBase {
+    private static final long DRAIN_TIMEOUT = 10;
     private Telephony2gUpdater mTelephony2gUpdater;
+    private SubscriptionManager.OnSubscriptionsChangedListener mChangedListener;
     private Executor mExecutor;
+    private CountDownLatch mLatch;
 
     private UserManager mMockUserManager;
     private TelephonyManager mMockTelephonyManager;
     private SubscriptionManager mMockSubscriptionManager;
 
+    // Set up to be returned from mMockSubscriptionManager.getCompleteActiveSubscriptionInfoList()
+    // Updates will be reflected in subsequent calls to the mock method.
+    private List<SubscriptionInfo> mCurrentSubscriptions;
+
     // 2G Bitmask is 0b10000000_01001011
     private static final long BASE_NETWORK = 0b11111111_11111111;
     private static final long EXPECTED_DISABLED = 0b01111111_10110100;
@@ -65,89 +80,247 @@
         mMockUserManager = mContext.getSystemService(UserManager.class);
         mMockSubscriptionManager = mContext.getSystemService(SubscriptionManager.class);
 
+        mCurrentSubscriptions = new ArrayList<>();
+        setupMutableSubscriptionInfoMock();
+
         mExecutor = Executors.newSingleThreadExecutor();
-        mTelephony2gUpdater = new Telephony2gUpdater(mExecutor,
-                getTestContext(), BASE_NETWORK);
+        mTelephony2gUpdater = new Telephony2gUpdater(mExecutor, getTestContext(), BASE_NETWORK);
+        mTelephony2gUpdater.init();
+        ArgumentCaptor<SubscriptionManager.OnSubscriptionsChangedListener> argument =
+                ArgumentCaptor.forClass(SubscriptionManager.OnSubscriptionsChangedListener.class);
+        verify(mMockSubscriptionManager).addOnSubscriptionsChangedListener(any(Executor.class),
+                argument.capture());
+        mChangedListener = argument.getValue();
     }
 
     @Test
-    public void handleUserRestrictionsChanged_noSubscriptions_noAllowedNetworksChanged() {
-        when(mMockSubscriptionManager.getAvailableSubscriptionInfoList()).thenReturn(
-                new ArrayList<>());
-        mTelephony2gUpdater.handleUserRestrictionsChanged(getTestContext());
+    public void onSubscriptionsChanged_noSubscriptions_noAllowedNetworksChanged() {
+        triggerOnSubscriptionChangedAndWait();
         verify(mMockTelephonyManager, never()).setAllowedNetworkTypesForReason(anyInt(), anyInt());
     }
 
     @Test
-    public void handleUserRestrictionsChanged_nullSubscriptions_noAllowedNetworksChanged() {
-        when(mMockSubscriptionManager.getAvailableSubscriptionInfoList()).thenReturn(null);
-        mTelephony2gUpdater.handleUserRestrictionsChanged(getTestContext());
-        verify(mMockTelephonyManager, never()).setAllowedNetworkTypesForReason(anyInt(), anyInt());
-    }
+    public void onSubscriptionsChanged_oneSubscription_allowedNetworksUpdated() {
+        TelephonyManager tmSubscription1 = addSubscriptionAndGetMock(1001);
+        triggerOnSubscriptionChangedAndWait();
 
-    @Test
-    public void handleUserRestrictionsChanged_oneSubscription_allowedNetworksUpdated() {
-        TelephonyManager tmSubscription1 = mock(TelephonyManager.class);
-        when(mMockSubscriptionManager.getAvailableSubscriptionInfoList()).thenReturn(
-                Collections.singletonList(getSubInfo(1)));
-        when(mMockTelephonyManager.createForSubscriptionId(1)).thenReturn(tmSubscription1);
-        when(mMockUserManager.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G)).thenReturn(
-                true);
-
-        mTelephony2gUpdater.handleUserRestrictionsChanged(getTestContext());
-
-        System.out.println(TelephonyManager.convertNetworkTypeBitmaskToString(11L));
         verify(tmSubscription1, times(1)).setAllowedNetworkTypesForReason(
-                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
     }
 
     @Test
-    public void handleUserRestrictionsChanged_manySubscriptionsDisallow2g_allowedNetworkUpdated() {
-
-        // Two subscriptions are available
-        when(mMockSubscriptionManager.getAvailableSubscriptionInfoList()).thenReturn(
-                Arrays.asList(getSubInfo(1), getSubInfo(2)));
-        TelephonyManager tmSubscription1 = mock(TelephonyManager.class);
-        TelephonyManager tmSubscription2 = mock(TelephonyManager.class);
-        when(mMockTelephonyManager.createForSubscriptionId(1)).thenReturn(tmSubscription1);
-        when(mMockTelephonyManager.createForSubscriptionId(2)).thenReturn(tmSubscription2);
+    public void onSubscriptionsChanged_manySubscriptionsDisallow2g_allowedNetworkUpdated() {
         // 2g is disallowed
         when(mMockUserManager.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G)).thenReturn(
                 true);
+        triggerBroadcastReceiverAndWait();
 
-        mTelephony2gUpdater.handleUserRestrictionsChanged(getTestContext());
+        TelephonyManager tmSubscription1 = addSubscriptionAndGetMock(1001);
+        TelephonyManager tmSubscription2 = addSubscriptionAndGetMock(1002);
+
+        triggerOnSubscriptionChangedAndWait();
 
         verify(tmSubscription1, times(1)).setAllowedNetworkTypesForReason(
                 TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+        verify(tmSubscription2, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+    }
+
+    @Test
+    public void onSubscriptionsChanged_noNewSubscriptions_noAllowedNetworksChanged() {
+        when(mMockUserManager.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G)).thenReturn(
+                true);
+        triggerBroadcastReceiverAndWait();
+
+        TelephonyManager tmSubscription1 = addSubscriptionAndGetMock(1001);
+
+        triggerOnSubscriptionChangedAndWait();
+        triggerOnSubscriptionChangedAndWait();
+
+        // subscriptions were updated twice, but we have no new subIds so we only expect one update
         verify(tmSubscription1, times(1)).setAllowedNetworkTypesForReason(
                 TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
     }
 
     @Test
-    public void handleUserRestrictionsChanged_manySubscriptionsAllow2g_allowedNetworkUpdated() {
+    public void onSubscriptionsChanged_removeSubscription_noAdditionalNetworkChanges() {
+        // We start with 2 subIds
+        TelephonyManager tmSubscription1 = addSubscriptionAndGetMock(1001);
+        TelephonyManager tmSubscription2 = addSubscriptionAndGetMock(1002);
 
-        // Two subscriptions are available
-        when(mMockSubscriptionManager.getAvailableSubscriptionInfoList()).thenReturn(
-                Arrays.asList(getSubInfo(1), getSubInfo(2)));
-        TelephonyManager tmSubscription1 = mock(TelephonyManager.class);
-        TelephonyManager tmSubscription2 = mock(TelephonyManager.class);
-        when(mMockTelephonyManager.createForSubscriptionId(1)).thenReturn(tmSubscription1);
-        when(mMockTelephonyManager.createForSubscriptionId(2)).thenReturn(tmSubscription2);
+        triggerOnSubscriptionChangedAndWait();
 
-        // 2g is allowed
+        // 2g is still enabled since the default is to not set the user restriction
+        verify(tmSubscription1, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+        verify(tmSubscription2, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+
+
+        mCurrentSubscriptions.remove(1);
+        triggerOnSubscriptionChangedAndWait();
+
+        // Subscriptions have changed, but we've only removed a subscription so there should be no
+        // extra updates to allowed network types
+        verify(tmSubscription1, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+        verify(tmSubscription2, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+    }
+
+    @Test
+    public void onSubscriptionsChanged_removeSubscriptionAndReAdd() {
+        when(mMockUserManager.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G)).thenReturn(
+                true);
+        triggerBroadcastReceiverAndWait();
+
+        TelephonyManager tmSubscription1 = addSubscriptionAndGetMock(1001);
+        triggerOnSubscriptionChangedAndWait();
+        mCurrentSubscriptions.remove(0);
+        triggerOnSubscriptionChangedAndWait();
+        mCurrentSubscriptions.add(getSubInfo(1001));
+        triggerOnSubscriptionChangedAndWait();
+
+        // subscriptions were updated thrice, but one of those updates removed a subscription
+        // such that the sub list was empty, so we only expect an update on the first and last
+        // updates.
+        verify(tmSubscription1, times(2)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+    }
+
+    @Test
+    public void onSubscriptionsChanged_addSubscription_updateAllowedNetworks() {
+        // We start with 2 subIds and update subscriptions
+        TelephonyManager tmSubscription1 = addSubscriptionAndGetMock(1001);
+        TelephonyManager tmSubscription2 = addSubscriptionAndGetMock(1002);
+        triggerOnSubscriptionChangedAndWait();
+
+        // Then add a subId and update subscriptions again
+        TelephonyManager tmSubscription3 = addSubscriptionAndGetMock(1003);
+        triggerOnSubscriptionChangedAndWait();
+
+        // we only need to update the new subscription
+        verify(tmSubscription1, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+        verify(tmSubscription2, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+        verify(tmSubscription3, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+    }
+
+    @Test
+    public void onUserRestrictionUnchanged_noChangeToRestriction_noAllowedNetworksUpdated() {
+        TelephonyManager tmSubscription = addSubscriptionAndGetMock(1001);
+        triggerOnSubscriptionChangedAndWait();
+        // precondition: we've updated allowed networks to the default (2g enabled)
+        verify(tmSubscription, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+        verify(tmSubscription, never()).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+
+        when(mMockUserManager.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G)).thenReturn(
+                true);
+        triggerBroadcastReceiverAndWait();
+        triggerBroadcastReceiverAndWait();
+
+        // expect we only updated once even though we got two broadcasts for user restriction
+        // updates
+        verify(tmSubscription, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+        // extra check to ensure we haven't also somehow updated back to enabled along the way
+        verify(tmSubscription, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+    }
+
+    @Test
+    public void onUserRestrictionChanged_restrictionChanged_allowedNetworksUpdated() {
+        // precondition: we've updated allowed networks to the default (2g enabled)
+        TelephonyManager tmSubscription = addSubscriptionAndGetMock(1001);
+        triggerOnSubscriptionChangedAndWait();
+        verify(tmSubscription, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+        verify(tmSubscription, never()).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+
+        // update the user restriction to disallow 2g
+        reset(tmSubscription);
+        when(mMockUserManager.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G)).thenReturn(
+                true);
+        triggerBroadcastReceiverAndWait();
+        verify(tmSubscription, never()).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+        verify(tmSubscription, times(1)).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+
+
+        // update the user restriction to allow 2g again
+        reset(tmSubscription);
         when(mMockUserManager.hasUserRestriction(UserManager.DISALLOW_CELLULAR_2G)).thenReturn(
                 false);
-
-        mTelephony2gUpdater.handleUserRestrictionsChanged(getTestContext());
-
-        verify(tmSubscription1, times(1)).setAllowedNetworkTypesForReason(
+        triggerBroadcastReceiverAndWait();
+        verify(tmSubscription, times(1)).setAllowedNetworkTypesForReason(
                 TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
-        verify(tmSubscription1, times(1)).setAllowedNetworkTypesForReason(
-                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_ENABLED);
+        verify(tmSubscription, never()).setAllowedNetworkTypesForReason(
+                TelephonyManager.ALLOWED_NETWORK_TYPES_REASON_USER_RESTRICTIONS, EXPECTED_DISABLED);
+
     }
 
     private SubscriptionInfo getSubInfo(int id) {
         return new SubscriptionInfo(id, "890126042XXXXXXXXXXX", 0, "T-mobile", "T-mobile", 0, 255,
                 "12345", 0, null, "310", "260", "156", false, null, null);
     }
+
+    private void triggerOnSubscriptionChangedAndWait() {
+        mExecutor.execute(() -> mChangedListener.onSubscriptionsChanged());
+        drainSingleThreadedExecutor();
+    }
+
+    private void triggerBroadcastReceiverAndWait() {
+        mTelephony2gUpdater.onReceive(mContext, new Intent());
+        drainSingleThreadedExecutor();
+    }
+
+    /**
+     * Wait for all tasks on executor up to the point of invocation to drain, then return.
+     *
+     * This helper takes advantage of the fact that we're using an immutable single threaded
+     * executor that guarantees tasks are executed in the order they are enqueued. It enqueues a
+     * task that decrements a latch and then waits on that task to finish. By definition, once the
+     * test task finishes, all previously enqueued tasks will have also completed.
+     */
+    private void drainSingleThreadedExecutor() {
+        resetExecutorLatch();
+        mExecutor.execute(() -> mLatch.countDown());
+        try {
+            mLatch.await(DRAIN_TIMEOUT, TimeUnit.SECONDS);
+        } catch (InterruptedException e) {
+            fail(e.getMessage());
+        }
+    }
+
+    private void resetExecutorLatch() {
+        mLatch = new CountDownLatch(1);
+    }
+
+    /**
+     * Helper that allows you to update subInfo and have that change reflected on subsequent calls
+     * to {@link SubscriptionManager#getCompleteActiveSubscriptionInfoList()}
+     */
+    private void setupMutableSubscriptionInfoMock() {
+        var answer = new Answer<List<SubscriptionInfo>>() {
+            @Override
+            public List<SubscriptionInfo> answer(InvocationOnMock invocation) throws Throwable {
+                return mCurrentSubscriptions;
+            }
+        };
+        when(mMockSubscriptionManager.getCompleteActiveSubscriptionInfoList()).thenAnswer(answer);
+    }
+
+    private TelephonyManager addSubscriptionAndGetMock(int subId) {
+        mCurrentSubscriptions.add(getSubInfo(subId));
+        TelephonyManager tmSubscription = mock(TelephonyManager.class);
+        when(mMockTelephonyManager.createForSubscriptionId(subId)).thenReturn(tmSubscription);
+        return tmSubscription;
+    }
+
 }