Update tests for VPN data stall recovery

Data stall recovery in VPN does not rely on the data stall
event now. Update the tests accordingly.

Bug: 289746779
Test: atest FrameworksNetTests
Change-Id: I4d2d96fe2f76706512325b6653b3e7cced10f207
diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index dc50773..b2f5cf8 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -2768,23 +2768,31 @@
                 new PersistableBundle());
     }
 
-    private void verifyMobikeTriggered(List<Network> expected) {
+    private void verifyMobikeTriggered(List<Network> expected, int retryIndex) {
+        // Verify retry is scheduled
+        final long expectedDelaySec = mTestDeps.getValidationFailRecoverySeconds(retryIndex);
+        final ArgumentCaptor<Long> delayCaptor = ArgumentCaptor.forClass(Long.class);
+        verify(mExecutor, times(retryIndex + 1)).schedule(
+                any(Runnable.class), delayCaptor.capture(), eq(TimeUnit.SECONDS));
+        final List<Long> delays = delayCaptor.getAllValues();
+        assertEquals(expectedDelaySec, (long) delays.get(delays.size() - 1));
+
         final ArgumentCaptor<Network> networkCaptor = ArgumentCaptor.forClass(Network.class);
-        verify(mIkeSessionWrapper).setNetwork(networkCaptor.capture(),
-                anyInt() /* ipVersion */, anyInt() /* encapType */, anyInt() /* keepaliveDelay */);
+        // TODO: Make the timeout shorter if real timeout will be used
+        verify(mIkeSessionWrapper, timeout(TEST_TIMEOUT_MS + expectedDelaySec * 1000))
+                .setNetwork(networkCaptor.capture(), anyInt() /* ipVersion */,
+                        anyInt() /* encapType */, anyInt() /* keepaliveDelay */);
         assertEquals(expected, Collections.singletonList(networkCaptor.getValue()));
     }
 
     @Test
     public void testDataStallInIkev2VpnMobikeDisabled() throws Exception {
-        verifySetupPlatformVpn(
+        final PlatformVpnSnapshot vpnSnapShot = verifySetupPlatformVpn(
                 createIkeConfig(createIkeConnectInfo(), false /* isMobikeEnabled */));
 
         doReturn(TEST_NETWORK).when(mMockNetworkAgent).getNetwork();
-        final ConnectivityDiagnosticsCallback connectivityDiagCallback =
-                getConnectivityDiagCallback();
-        final DataStallReport report = createDataStallReport();
-        connectivityDiagCallback.onDataStallSuspected(report);
+        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
+                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
 
         // Should not trigger MOBIKE if MOBIKE is not enabled
         verify(mIkeSessionWrapper, never()).setNetwork(any() /* network */,
@@ -2797,19 +2805,11 @@
                 createIkeConfig(createIkeConnectInfo(), true /* isMobikeEnabled */));
 
         doReturn(TEST_NETWORK).when(mMockNetworkAgent).getNetwork();
-        final ConnectivityDiagnosticsCallback connectivityDiagCallback =
-                getConnectivityDiagCallback();
-        final DataStallReport report = createDataStallReport();
-        connectivityDiagCallback.onDataStallSuspected(report);
-
+        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
+                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
         // Verify MOBIKE is triggered
-        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks());
-
-        // Expect to skip other data stall event if MOBIKE was started.
-        reset(mIkeSessionWrapper);
-        connectivityDiagCallback.onDataStallSuspected(report);
-        verify(mIkeSessionWrapper, never()).setNetwork(any() /* network */,
-                anyInt() /* ipVersion */, anyInt() /* encapType */, anyInt() /* keepaliveDelay */);
+        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks(),
+                0 /* retryIndex */);
 
         reset(mIkev2SessionCreator);
 
@@ -2819,14 +2819,6 @@
                 NetworkAgent.VALIDATION_STATUS_VALID);
         verify(mIkev2SessionCreator, never()).createIkeSession(
                 any(), any(), any(), any(), any(), any());
-
-        // Send invalid result to verify no ike session reset since the data stall suspected
-        // variables(timer counter and boolean) was reset.
-        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
-                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
-        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
-        verify(mIkev2SessionCreator, never()).createIkeSession(
-                any(), any(), any(), any(), any(), any());
     }
 
     @Test
@@ -2834,31 +2826,46 @@
         final PlatformVpnSnapshot vpnSnapShot = verifySetupPlatformVpn(
                 createIkeConfig(createIkeConnectInfo(), true /* isMobikeEnabled */));
 
-        final ConnectivityDiagnosticsCallback connectivityDiagCallback =
-                getConnectivityDiagCallback();
-
+        int retry = 0;
         doReturn(TEST_NETWORK).when(mMockNetworkAgent).getNetwork();
-        final DataStallReport report = createDataStallReport();
-        connectivityDiagCallback.onDataStallSuspected(report);
-
-        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks());
+        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
+                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
+        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks(),
+                retry++);
 
         reset(mIkev2SessionCreator);
 
+        // Second validation status update.
+        ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
+                NetworkAgent.VALIDATION_STATUS_NOT_VALID);
+        verifyMobikeTriggered(vpnSnapShot.vpn.mNetworkCapabilities.getUnderlyingNetworks(),
+                retry++);
+
+        // Use real delay to verify reset session will not be performed if there is an existing
+        // recovery for resetting the session.
+        mExecutor.delayMs = TestExecutor.REAL_DELAY;
+        mExecutor.executeDirect = true;
         // Send validation status update should result in ike session reset.
         ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
                 NetworkAgent.VALIDATION_STATUS_NOT_VALID);
 
-        // Verify reset is scheduled and run.
-        verify(mExecutor, atLeastOnce()).schedule(any(Runnable.class), anyLong(), any());
+        // Verify session reset is scheduled
+        long expectedDelay = mTestDeps.getValidationFailRecoverySeconds(retry++);
+        final ArgumentCaptor<Long> delayCaptor = ArgumentCaptor.forClass(Long.class);
+        verify(mExecutor, times(retry)).schedule(any(Runnable.class), delayCaptor.capture(),
+                eq(TimeUnit.SECONDS));
+        final List<Long> delays = delayCaptor.getAllValues();
+        assertEquals(expectedDelay, (long) delays.get(delays.size() - 1));
 
         // Another invalid status reported should not trigger other scheduled recovery.
-        reset(mExecutor);
+        expectedDelay = mTestDeps.getValidationFailRecoverySeconds(retry++);
         ((Vpn.IkeV2VpnRunner) vpnSnapShot.vpn.mVpnRunner).onValidationStatus(
                 NetworkAgent.VALIDATION_STATUS_NOT_VALID);
-        verify(mExecutor, never()).schedule(any(Runnable.class), anyLong(), any());
+        verify(mExecutor, never()).schedule(
+                any(Runnable.class), eq(expectedDelay), eq(TimeUnit.SECONDS));
 
-        verify(mIkev2SessionCreator, timeout(TEST_TIMEOUT_MS))
+        // Verify that session being reset
+        verify(mIkev2SessionCreator, timeout(TEST_TIMEOUT_MS + expectedDelay * 1000))
                 .createIkeSession(any(), any(), any(), any(), any(), any());
     }
 
@@ -3137,6 +3144,12 @@
         }
 
         @Override
+        public long getValidationFailRecoverySeconds(int retryCount) {
+            // Simply return retryCount as the delay seconds for retrying.
+            return retryCount;
+        }
+
+        @Override
         public ScheduledThreadPoolExecutor newScheduledThreadPoolExecutor() {
             return mExecutor;
         }