Adjust tests for delay in restarting/migrating IKE

This patch only adjusts the existing tests for the change
in the companion patch. Actual tests will come as a followup
because this patch is already big enough.

Test: VpnTest
Bug: 269715746
Change-Id: I65542a8f4151b4857f1b3758b2cae887bfbfe261
diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index dd9177ee..caf080c 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -308,14 +308,27 @@
     @Mock private SubscriptionManager mSubscriptionManager;
     @Mock private IpSecService mIpSecService;
     @Mock private VpnProfileStore mVpnProfileStore;
-    @Mock private ScheduledThreadPoolExecutor mExecutor;
-    @Mock private ScheduledFuture mScheduledFuture;
+    private final ScheduledThreadPoolExecutor mExecutor;
     @Mock DeviceIdleInternal mDeviceIdleInternal;
     private final VpnProfile mVpnProfile;
 
     private IpSecManager mIpSecManager;
     private TestDeps mTestDeps;
 
+    public static class TestExecutor extends ScheduledThreadPoolExecutor {
+        public TestExecutor() {
+            super(1);
+        }
+
+        // For the purposes of the test, run all scheduled tasks after 10ms to save
+        // execution time
+        @Override
+        public ScheduledFuture<?> schedule(final Runnable command, final long delay,
+                final TimeUnit unit) {
+            return super.schedule(command, 10, TimeUnit.MILLISECONDS);
+        }
+    }
+
     public VpnTest() throws Exception {
         // Build an actual VPN profile that is capable of being converted to and from an
         // Ikev2VpnProfile
@@ -323,6 +336,7 @@
                 new Ikev2VpnProfile.Builder(TEST_VPN_SERVER, TEST_VPN_IDENTITY);
         builder.setAuthPsk(TEST_VPN_PSK);
         builder.setBypassable(true /* isBypassable */);
+        mExecutor = spy(new TestExecutor());
         mVpnProfile = builder.build().toVpnProfile();
     }
 
@@ -388,7 +402,6 @@
 
         // Set up mIkev2SessionCreator and mExecutor
         resetIkev2SessionCreator(mIkeSessionWrapper);
-        resetExecutor(mScheduledFuture);
     }
 
     private void resetIkev2SessionCreator(Vpn.IkeSessionWrapper ikeSession) {
@@ -397,18 +410,6 @@
                 .thenReturn(ikeSession);
     }
 
-    private void resetExecutor(ScheduledFuture scheduledFuture) {
-        doAnswer(
-                (invocation) -> {
-                    ((Runnable) invocation.getArgument(0)).run();
-                    return null;
-                })
-            .when(mExecutor)
-            .execute(any());
-        when(mExecutor.schedule(
-                any(Runnable.class), anyLong(), any())).thenReturn(mScheduledFuture);
-    }
-
     @After
     public void tearDown() throws Exception {
         doReturn(PERMISSION_DENIED).when(mContext).checkCallingOrSelfPermission(any());
@@ -524,9 +525,9 @@
     }
 
     private void verifyPowerSaveTempWhitelistApp(String packageName) {
-        verify(mDeviceIdleInternal).addPowerSaveTempWhitelistApp(anyInt(), eq(packageName),
-                anyLong(), anyInt(), eq(false), eq(PowerWhitelistManager.REASON_VPN),
-                eq("VpnManager event"));
+        verify(mDeviceIdleInternal, timeout(TEST_TIMEOUT_MS)).addPowerSaveTempWhitelistApp(
+                anyInt(), eq(packageName), anyLong(), anyInt(), eq(false),
+                eq(PowerWhitelistManager.REASON_VPN), eq("VpnManager event"));
     }
 
     @Test
@@ -1574,10 +1575,7 @@
         // same process with the real case.
         if (errorCode == VpnManager.ERROR_CODE_NETWORK_LOST) {
             cb.onLost(TEST_NETWORK);
-            final ArgumentCaptor<Runnable> runnableCaptor =
-                    ArgumentCaptor.forClass(Runnable.class);
-            verify(mExecutor).schedule(runnableCaptor.capture(), anyLong(), any());
-            runnableCaptor.getValue().run();
+            verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
         } else {
             final IkeSessionCallback ikeCb = captor.getValue();
             ikeCb.onClosedWithException(exception);
@@ -1602,25 +1600,23 @@
     }
 
     private IkeSessionCallback verifyRetryAndGetNewIkeCb(int retryIndex) {
-        final ArgumentCaptor<Runnable> runnableCaptor =
-                ArgumentCaptor.forClass(Runnable.class);
         final ArgumentCaptor<IkeSessionCallback> ikeCbCaptor =
                 ArgumentCaptor.forClass(IkeSessionCallback.class);
 
         // Verify retry is scheduled
-        final long expectedDelay = mTestDeps.getNextRetryDelaySeconds(retryIndex);
-        verify(mExecutor).schedule(runnableCaptor.capture(), eq(expectedDelay), any());
+        final long expectedDelay = mTestDeps.getNextRetryDelaySeconds(retryIndex) * 1000;
+        final ArgumentCaptor<Long> delayCaptor = ArgumentCaptor.forClass(Long.class);
+        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), delayCaptor.capture(),
+                eq(TimeUnit.MILLISECONDS));
+        final List<Long> delays = delayCaptor.getAllValues();
+        assertEquals(expectedDelay, (long) delays.get(delays.size() - 1));
 
-        // Mock the event of firing the retry task
-        runnableCaptor.getValue().run();
-
-        verify(mIkev2SessionCreator)
+        verify(mIkev2SessionCreator, timeout(TEST_TIMEOUT_MS + expectedDelay))
                 .createIkeSession(any(), any(), any(), any(), ikeCbCaptor.capture(), any());
 
         // Forget the mIkev2SessionCreator#createIkeSession call and mExecutor#schedule call
         // for the next retry verification
         resetIkev2SessionCreator(mIkeSessionWrapper);
-        resetExecutor(mScheduledFuture);
 
         return ikeCbCaptor.getValue();
     }
@@ -1884,6 +1880,8 @@
 
         vpn.startVpnProfile(TEST_VPN_PKG);
         final NetworkCallback nwCb = triggerOnAvailableAndGetCallback();
+        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
+        reset(mExecutor);
 
         // Mock the setup procedure by firing callbacks
         final Pair<IkeSessionCallback, ChildSessionCallback> cbPair =
@@ -2088,7 +2086,7 @@
         vpnSnapShot.nwCb.onCapabilitiesChanged(
                 TEST_NETWORK_2, new NetworkCapabilities.Builder().build());
         // Verify MOBIKE is triggered
-        verify(mIkeSessionWrapper).setNetwork(TEST_NETWORK_2,
+        verify(mIkeSessionWrapper, timeout(TEST_TIMEOUT_MS)).setNetwork(TEST_NETWORK_2,
                 expectedIpVersion, expectedEncapType, expectedKeepalive);
 
         vpnSnapShot.vpn.mVpnRunner.exitVpnRunner();
@@ -2247,7 +2245,7 @@
         reset(mIkeSessionWrapper);
         mockCarrierConfig(TEST_SUB_ID, simState, TEST_KEEPALIVE_TIMER, preferredIpProto);
         vpnSnapShot.nwCb.onCapabilitiesChanged(TEST_NETWORK_2, nc);
-        verify(mIkeSessionWrapper).setNetwork(TEST_NETWORK_2,
+        verify(mIkeSessionWrapper, timeout(TEST_TIMEOUT_MS)).setNetwork(TEST_NETWORK_2,
                 expectedIpVersion, expectedEncapType, expectedKeepaliveTimer);
         if (expectedReadFromCarrierConfig) {
             final ArgumentCaptor<NetworkCapabilities> ncCaptor =
@@ -2296,17 +2294,16 @@
 
         // Mock network loss and verify a cleanup task is scheduled
         vpnSnapShot.nwCb.onLost(TEST_NETWORK);
-        verify(mExecutor).schedule(any(Runnable.class), anyLong(), any());
+        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
 
         // Mock new network comes up and the cleanup task is cancelled
         vpnSnapShot.nwCb.onAvailable(TEST_NETWORK_2);
-        verify(mScheduledFuture).cancel(anyBoolean());
         verify(mIkeSessionWrapper, never()).setNetwork(any(), anyInt(), anyInt(), anyInt());
 
         vpnSnapShot.nwCb.onCapabilitiesChanged(TEST_NETWORK_2,
                 new NetworkCapabilities.Builder().build());
         // Verify MOBIKE is triggered
-        verify(mIkeSessionWrapper).setNetwork(eq(TEST_NETWORK_2),
+        verify(mIkeSessionWrapper, timeout(TEST_TIMEOUT_MS)).setNetwork(eq(TEST_NETWORK_2),
                 eq(ESP_IP_VERSION_AUTO) /* ipVersion */,
                 eq(ESP_ENCAP_TYPE_AUTO) /* encapType */,
                 eq(DEFAULT_UDP_PORT_4500_NAT_TIMEOUT_SEC_INT) /* keepaliveDelay */);
@@ -2405,7 +2402,7 @@
         vpnSnapShot.nwCb.onCapabilitiesChanged(
                 TEST_NETWORK_2, new NetworkCapabilities.Builder().build());
         // Verify the old IKE Session is killed
-        verify(mIkeSessionWrapper).kill();
+        verify(mIkeSessionWrapper, timeout(TEST_TIMEOUT_MS)).kill();
 
         // Capture callbacks of the new IKE Session
         final Pair<IkeSessionCallback, ChildSessionCallback> cbPair =
@@ -2437,19 +2434,16 @@
         // Forget the #sendLinkProperties during first setup.
         reset(mMockNetworkAgent);
 
-        final ArgumentCaptor<Runnable> runnableCaptor =
-                ArgumentCaptor.forClass(Runnable.class);
-
         // Mock network loss
         vpnSnapShot.nwCb.onLost(TEST_NETWORK);
 
         // Mock the grace period expires
-        verify(mExecutor).schedule(runnableCaptor.capture(), anyLong(), any());
-        runnableCaptor.getValue().run();
+        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
 
         final ArgumentCaptor<LinkProperties> lpCaptor =
                 ArgumentCaptor.forClass(LinkProperties.class);
-        verify(mMockNetworkAgent).doSendLinkProperties(lpCaptor.capture());
+        verify(mMockNetworkAgent, timeout(TEST_TIMEOUT_MS))
+                .doSendLinkProperties(lpCaptor.capture());
         final LinkProperties lp = lpCaptor.getValue();
 
         assertNull(lp.getInterfaceName());
@@ -2547,9 +2541,7 @@
         // variables(timer counter and boolean) was reset.
         ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
                 NetworkAgent.VALIDATION_STATUS_NOT_VALID);
-        final ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class);
-        verify(mExecutor).schedule(runnableCaptor.capture(), anyLong(), any());
-        runnableCaptor.getValue().run();
+        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
         verify(mIkev2SessionCreator, never()).createIkeSession(
                 any(), any(), any(), any(), any(), any());
     }
@@ -2575,17 +2567,16 @@
                 NetworkAgent.VALIDATION_STATUS_NOT_VALID);
 
         // Verify reset is scheduled and run.
-        final ArgumentCaptor<Runnable> runnableCaptor = ArgumentCaptor.forClass(Runnable.class);
-        verify(mExecutor).schedule(runnableCaptor.capture(), anyLong(), any());
+        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
 
         // Another invalid status reported should not trigger other scheduled recovery.
         reset(mExecutor);
         ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
                 NetworkAgent.VALIDATION_STATUS_NOT_VALID);
-        verify(mExecutor, never()).schedule(runnableCaptor.capture(), anyLong(), any());
+        verify(mExecutor, never()).schedule(any(Runnable.class), anyLong(), any());
 
-        runnableCaptor.getValue().run();
-        verify(mIkev2SessionCreator).createIkeSession(any(), any(), any(), any(), any(), any());
+        verify(mIkev2SessionCreator, timeout(TEST_TIMEOUT_MS))
+                .createIkeSession(any(), any(), any(), any(), any(), any());
     }
 
     @Test