2 files changed, 51 insertions, 0 deletions
diff --git a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionRetryTimeoutStateTest.java b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionRetryTimeoutStateTest.java
index a88f112f4502..69407657b3c4 100644
--- a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionRetryTimeoutStateTest.java
+++ b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionRetryTimeoutStateTest.java
@@ -16,10 +16,16 @@
 
 package com.android.server.vcn;
 
+import static com.android.server.vcn.VcnGatewayConnection.VcnNetworkAgent;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 
 import androidx.test.filters.SmallTest;
 import androidx.test.runner.AndroidJUnit4;
@@ -33,16 +39,20 @@ import org.junit.runner.RunWith;
 @SmallTest
 public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnectionTestBase {
     private long mFirstRetryInterval;
+    private VcnNetworkAgent mNetworkAgent;
 
     @Before
     public void setUp() throws Exception {
         super.setUp();
 
         mFirstRetryInterval = mConfig.getRetryIntervalsMillis()[0];
+        mNetworkAgent = mock(VcnNetworkAgent.class);
 
         mGatewayConnection.setUnderlyingNetwork(TEST_UNDERLYING_NETWORK_RECORD_1);
         mGatewayConnection.transitionTo(mGatewayConnection.mRetryTimeoutState);
         mTestLooper.dispatchAll();
+
+        mGatewayConnection.setNetworkAgent(mNetworkAgent);
     }
 
     @Test
@@ -54,6 +64,9 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
 
         assertEquals(mGatewayConnection.mConnectingState, mGatewayConnection.getCurrentState());
         verifyRetryTimeoutAlarmAndGetCallback(mFirstRetryInterval, true /* expectCanceled */);
+
+        assertNotNull(mGatewayConnection.getNetworkAgent());
+        verify(mNetworkAgent, never()).unregister();
     }
 
     @Test
@@ -65,6 +78,9 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
 
         assertEquals(mGatewayConnection.mRetryTimeoutState, mGatewayConnection.getCurrentState());
         verifyRetryTimeoutAlarmAndGetCallback(mFirstRetryInterval, false /* expectCanceled */);
+
+        assertNotNull(mGatewayConnection.getNetworkAgent());
+        verify(mNetworkAgent, never()).unregister();
     }
 
     @Test
@@ -76,6 +92,9 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
 
         assertEquals(mGatewayConnection.mDisconnectedState, mGatewayConnection.getCurrentState());
         verifyRetryTimeoutAlarmAndGetCallback(mFirstRetryInterval, true /* expectCanceled */);
+
+        assertNull(mGatewayConnection.getNetworkAgent());
+        verify(mNetworkAgent).unregister();
     }
 
     @Test
@@ -93,6 +112,9 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
 
         assertEquals(mGatewayConnection.mConnectingState, mGatewayConnection.getCurrentState());
         verifyRetryTimeoutAlarmAndGetCallback(mFirstRetryInterval, true /* expectCanceled */);
+
+        assertNotNull(mGatewayConnection.getNetworkAgent());
+        verify(mNetworkAgent, never()).unregister();
     }
 
     @Test
@@ -108,6 +130,9 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
 
         assertNull(mGatewayConnection.getCurrentState());
         assertTrue(mGatewayConnection.isQuitting());
+
+        assertNull(mGatewayConnection.getNetworkAgent());
+        verify(mNetworkAgent).unregister();
     }
 
     @Test
@@ -117,5 +142,8 @@ public class VcnGatewayConnectionRetryTimeoutStateTest extends VcnGatewayConnect
 
         assertEquals(mGatewayConnection.mDisconnectedState, mGatewayConnection.getCurrentState());
         assertFalse(mGatewayConnection.isQuitting());
+
+        assertNull(mGatewayConnection.getNetworkAgent());
+        verify(mNetworkAgent).unregister();
     }
 }
diff --git a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java
index a4f95e03e9bd..83610e0b7a67 100644
--- a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java
+++ b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTest.java
@@ -24,8 +24,12 @@ import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VCN_MANAGED;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 
+import static com.android.server.vcn.VcnGatewayConnection.VcnIkeSession;
+import static com.android.server.vcn.VcnGatewayConnection.VcnNetworkAgent;
+
 import static org.junit.Assert.assertArrayEquals;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Matchers.eq;
@@ -191,4 +195,23 @@ public class VcnGatewayConnectionTest extends VcnGatewayConnectionTestBase {
 
         verify(mDisconnectRequestAlarm).cancel();
     }
+
+    @Test
+    public void testQuittingCleansUpPersistentState() {
+        final VcnIkeSession vcnIkeSession = mock(VcnIkeSession.class);
+        final VcnNetworkAgent vcnNetworkAgent = mock(VcnNetworkAgent.class);
+
+        mGatewayConnection.setIkeSession(vcnIkeSession);
+        mGatewayConnection.setNetworkAgent(vcnNetworkAgent);
+
+        mGatewayConnection.quitNow();
+        mTestLooper.dispatchAll();
+
+        assertNull(mGatewayConnection.getIkeSession());
+        verify(vcnIkeSession).kill();
+        assertNull(mGatewayConnection.getNetworkAgent());
+        verify(vcnNetworkAgent).unregister();
+
+        verifyWakeLockReleased();
+    }
 }