diff --git a/tests/unit/java/com/android/server/connectivity/VpnTest.java b/tests/unit/java/com/android/server/connectivity/VpnTest.java
index 2ca728d..1f965d9 100644
--- a/tests/unit/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/unit/java/com/android/server/connectivity/VpnTest.java
@@ -2039,7 +2039,10 @@
                         false /* mtuSupportsIpv6 */);
         // Simulate a new network coming up
         vpnSnapShot.nwCb.onAvailable(TEST_NETWORK_2);
+        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(TEST_NETWORK_2,
                 expectedIpVersion, expectedEncapType, expectedKeepalive);
@@ -2124,11 +2127,8 @@
 
         // Simulate a new network coming up
         vpnSnapShot.nwCb.onAvailable(TEST_NETWORK_2);
-        // guessNattKeepaliveTimerForNetwork() will return AUTOMATIC_KEEPALIVE_DELAY_SECONDS
-        // because mUnderlyingNetworkCapabilities is not set at the moment.
-        // TODO: Fix it to prevent setNetwork() being triggered w/o valid network capabilities.
-        verify(mIkeSessionWrapper).setNetwork(TEST_NETWORK_2,
-                ESP_IP_VERSION_AUTO, ESP_ENCAP_TYPE_AUTO, AUTOMATIC_KEEPALIVE_DELAY_SECONDS);
+        // Migration will not be started until receiving network capabilities change.
+        verify(mIkeSessionWrapper, never()).setNetwork(any(), anyInt(), anyInt(), anyInt());
 
         reset(mIkeSessionWrapper);
         mockCarrierConfig(TEST_SUB_ID, TEST_KEEPALIVE_TIMER, simState);
@@ -2171,7 +2171,10 @@
         // 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),
                 eq(ESP_IP_VERSION_AUTO) /* ipVersion */,
@@ -2266,7 +2269,11 @@
         // Mock network switch
         vpnSnapShot.nwCb.onLost(TEST_NETWORK);
         vpnSnapShot.nwCb.onAvailable(TEST_NETWORK_2);
+        // The old IKE Session will not be killed until receiving network capabilities change.
+        verify(mIkeSessionWrapper, never()).kill();
 
+        vpnSnapShot.nwCb.onCapabilitiesChanged(
+                TEST_NETWORK_2, new NetworkCapabilities.Builder().build());
         // Verify the old IKE Session is killed
         verify(mIkeSessionWrapper).kill();