Create proxy NetworkAgent for testing

This change adds a proxy NetworkAgent to facilitate testing, adding the
ability to test final methods such as sendNetworkCapabilities() and
sendLinkProperties(), which would otherwise call a real implementation,
and throw a null pointer while trying to access internal fields.

Bug: 184579891
Bug: 184697651
Test: atest FrameworksVcnTests
Change-Id: I8f13159bd266bf488f095a13d6ffe444bec65f38
diff --git a/services/core/java/com/android/server/vcn/VcnGatewayConnection.java b/services/core/java/com/android/server/vcn/VcnGatewayConnection.java
index ddf8b46..8d24c3a 100644
--- a/services/core/java/com/android/server/vcn/VcnGatewayConnection.java
+++ b/services/core/java/com/android/server/vcn/VcnGatewayConnection.java
@@ -575,7 +575,7 @@
      * <p>Set in Connected state, always @NonNull in Connected, Migrating states, @Nullable
      * otherwise.
      */
-    private NetworkAgent mNetworkAgent;
+    private VcnNetworkAgent mNetworkAgent;
 
     @Nullable private WakeupMessage mTeardownTimeoutAlarm;
     @Nullable private WakeupMessage mDisconnectRequestAlarm;
@@ -1441,7 +1441,7 @@
     private abstract class ConnectedStateBase extends ActiveBaseState {
         protected void updateNetworkAgent(
                 @NonNull IpSecTunnelInterface tunnelIface,
-                @NonNull NetworkAgent agent,
+                @NonNull VcnNetworkAgent agent,
                 @NonNull VcnChildSessionConfiguration childConfig) {
             final NetworkCapabilities caps =
                     buildNetworkCapabilities(mConnectionConfig, mUnderlying);
@@ -1452,7 +1452,7 @@
             agent.sendLinkProperties(lp);
         }
 
-        protected NetworkAgent buildNetworkAgent(
+        protected VcnNetworkAgent buildNetworkAgent(
                 @NonNull IpSecTunnelInterface tunnelIface,
                 @NonNull VcnChildSessionConfiguration childConfig) {
             final NetworkCapabilities caps =
@@ -1464,7 +1464,7 @@
                             .setLegacyType(ConnectivityManager.TYPE_MOBILE)
                             .build();
 
-            final NetworkAgent agent =
+            final VcnNetworkAgent agent =
                     mDeps.newNetworkAgent(
                             mVcnContext,
                             TAG,
@@ -1993,12 +1993,12 @@
     }
 
     @VisibleForTesting(visibility = Visibility.PRIVATE)
-    NetworkAgent getNetworkAgent() {
+    VcnNetworkAgent getNetworkAgent() {
         return mNetworkAgent;
     }
 
     @VisibleForTesting(visibility = Visibility.PRIVATE)
-    void setNetworkAgent(@Nullable NetworkAgent networkAgent) {
+    void setNetworkAgent(@Nullable VcnNetworkAgent networkAgent) {
         mNetworkAgent = networkAgent;
     }
 
@@ -2086,8 +2086,8 @@
             return new WakeupMessage(vcnContext.getContext(), handler, tag, runnable);
         }
 
-        /** Builds a new NetworkAgent. */
-        public NetworkAgent newNetworkAgent(
+        /** Builds a new VcnNetworkAgent. */
+        public VcnNetworkAgent newNetworkAgent(
                 @NonNull VcnContext vcnContext,
                 @NonNull String tag,
                 @NonNull NetworkCapabilities caps,
@@ -2095,27 +2095,18 @@
                 @NonNull int score,
                 @NonNull NetworkAgentConfig nac,
                 @NonNull NetworkProvider provider,
-                @NonNull Consumer<NetworkAgent> networkUnwantedCallback,
+                @NonNull Consumer<VcnNetworkAgent> networkUnwantedCallback,
                 @NonNull Consumer<Integer> validationStatusCallback) {
-            return new NetworkAgent(
-                    vcnContext.getContext(),
-                    vcnContext.getLooper(),
+            return new VcnNetworkAgent(
+                    vcnContext,
                     tag,
                     caps,
                     lp,
                     score,
                     nac,
-                    provider) {
-                @Override
-                public void onNetworkUnwanted() {
-                    networkUnwantedCallback.accept(this);
-                }
-
-                @Override
-                public void onValidationStatus(int status, @Nullable Uri redirectUri) {
-                    validationStatusCallback.accept(status);
-                }
-            };
+                    provider,
+                    networkUnwantedCallback,
+                    validationStatusCallback);
         }
 
         /** Gets the elapsed real time since boot, in millis. */
@@ -2231,4 +2222,73 @@
             mImpl.release();
         }
     }
+
+    /** Proxy Implementation of NetworkAgent, used for testing. */
+    @VisibleForTesting(visibility = Visibility.PRIVATE)
+    public static class VcnNetworkAgent {
+        private final NetworkAgent mImpl;
+
+        public VcnNetworkAgent(
+                @NonNull VcnContext vcnContext,
+                @NonNull String tag,
+                @NonNull NetworkCapabilities caps,
+                @NonNull LinkProperties lp,
+                @NonNull int score,
+                @NonNull NetworkAgentConfig nac,
+                @NonNull NetworkProvider provider,
+                @NonNull Consumer<VcnNetworkAgent> networkUnwantedCallback,
+                @NonNull Consumer<Integer> validationStatusCallback) {
+            mImpl =
+                    new NetworkAgent(
+                            vcnContext.getContext(),
+                            vcnContext.getLooper(),
+                            tag,
+                            caps,
+                            lp,
+                            score,
+                            nac,
+                            provider) {
+                        @Override
+                        public void onNetworkUnwanted() {
+                            networkUnwantedCallback.accept(VcnNetworkAgent.this);
+                        }
+
+                        @Override
+                        public void onValidationStatus(int status, @Nullable Uri redirectUri) {
+                            validationStatusCallback.accept(status);
+                        }
+                    };
+        }
+
+        /** Registers the underlying NetworkAgent */
+        public void register() {
+            mImpl.register();
+        }
+
+        /** Marks the underlying NetworkAgent as connected */
+        public void markConnected() {
+            mImpl.markConnected();
+        }
+
+        /** Unregisters the underlying NetworkAgent */
+        public void unregister() {
+            mImpl.unregister();
+        }
+
+        /** Sends new NetworkCapabilities for the underlying NetworkAgent */
+        public void sendNetworkCapabilities(@NonNull NetworkCapabilities caps) {
+            mImpl.sendNetworkCapabilities(caps);
+        }
+
+        /** Sends new LinkProperties for the underlying NetworkAgent */
+        public void sendLinkProperties(@NonNull LinkProperties lp) {
+            mImpl.sendLinkProperties(lp);
+        }
+
+        /** Retrieves the Network for the underlying NetworkAgent */
+        @Nullable
+        public Network getNetwork() {
+            return mImpl.getNetwork();
+        }
+    }
 }
diff --git a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionConnectedStateTest.java b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionConnectedStateTest.java
index 5dcf4a8..cb7c96a 100644
--- a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionConnectedStateTest.java
+++ b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionConnectedStateTest.java
@@ -26,6 +26,7 @@
 
 import static com.android.server.vcn.VcnGatewayConnection.VcnChildSessionConfiguration;
 import static com.android.server.vcn.VcnGatewayConnection.VcnIkeSession;
+import static com.android.server.vcn.VcnGatewayConnection.VcnNetworkAgent;
 
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
@@ -73,13 +74,13 @@
 @SmallTest
 public class VcnGatewayConnectionConnectedStateTest extends VcnGatewayConnectionTestBase {
     private VcnIkeSession mIkeSession;
-    private NetworkAgent mNetworkAgent;
+    private VcnNetworkAgent mNetworkAgent;
 
     @Before
     public void setUp() throws Exception {
         super.setUp();
 
-        mNetworkAgent = mock(NetworkAgent.class);
+        mNetworkAgent = mock(VcnNetworkAgent.class);
         doReturn(mNetworkAgent)
                 .when(mDeps)
                 .newNetworkAgent(any(), any(), any(), any(), anyInt(), any(), any(), any(), any());
@@ -299,8 +300,9 @@
                 .removeAddressFromTunnelInterface(
                         eq(TEST_IPSEC_TUNNEL_RESOURCE_ID), eq(TEST_INTERNAL_ADDR), any());
 
-        // TODO(b/184579891): Also verify link properties updated and sent when sendLinkProperties
-        // is mockable
+        verify(mNetworkAgent).sendLinkProperties(argThat(
+                lp -> newInternalAddrs.equals(lp.getLinkAddresses())
+                        && Collections.singletonList(TEST_DNS_ADDR_2).equals(lp.getDnsServers())));
 
         // Verify that IpSecTunnelInterface only created once
         verify(mIpSecSvc).createTunnelInterface(any(), any(), any(), any(), any());
@@ -323,11 +325,11 @@
         assertFalse(mGatewayConnection.isInSafeMode());
     }
 
-    private Consumer<NetworkAgent> setupNetworkAndGetUnwantedCallback() {
+    private Consumer<VcnNetworkAgent> setupNetworkAndGetUnwantedCallback() {
         triggerChildOpened();
         mTestLooper.dispatchAll();
 
-        final ArgumentCaptor<Consumer<NetworkAgent>> unwantedCallbackCaptor =
+        final ArgumentCaptor<Consumer<VcnNetworkAgent>> unwantedCallbackCaptor =
                 ArgumentCaptor.forClass(Consumer.class);
         verify(mDeps)
                 .newNetworkAgent(
@@ -346,7 +348,7 @@
 
     @Test
     public void testUnwantedNetworkAgentTriggersTeardown() throws Exception {
-        final Consumer<NetworkAgent> unwantedCallback = setupNetworkAndGetUnwantedCallback();
+        final Consumer<VcnNetworkAgent> unwantedCallback = setupNetworkAndGetUnwantedCallback();
 
         unwantedCallback.accept(mNetworkAgent);
         mTestLooper.dispatchAll();
@@ -357,7 +359,7 @@
 
     @Test
     public void testUnwantedNetworkAgentWithDisconnectedNetworkAgent() throws Exception {
-        final Consumer<NetworkAgent> unwantedCallback = setupNetworkAndGetUnwantedCallback();
+        final Consumer<VcnNetworkAgent> unwantedCallback = setupNetworkAndGetUnwantedCallback();
 
         mGatewayConnection.setNetworkAgent(null);
         unwantedCallback.accept(mNetworkAgent);
@@ -371,8 +373,8 @@
 
     @Test
     public void testUnwantedNetworkAgentWithNewNetworkAgent() throws Exception {
-        final Consumer<NetworkAgent> unwantedCallback = setupNetworkAndGetUnwantedCallback();
-        final NetworkAgent testAgent = mock(NetworkAgent.class);
+        final Consumer<VcnNetworkAgent> unwantedCallback = setupNetworkAndGetUnwantedCallback();
+        final VcnNetworkAgent testAgent = mock(VcnNetworkAgent.class);
 
         mGatewayConnection.setNetworkAgent(testAgent);
         unwantedCallback.accept(mNetworkAgent);
diff --git a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
index c5ed8f6..b93347b 100644
--- a/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
+++ b/tests/vcn/java/com/android/server/vcn/VcnGatewayConnectionTestBase.java
@@ -18,6 +18,7 @@
 
 import static com.android.server.vcn.UnderlyingNetworkTracker.UnderlyingNetworkRecord;
 import static com.android.server.vcn.VcnGatewayConnection.VcnIkeSession;
+import static com.android.server.vcn.VcnGatewayConnection.VcnNetworkAgent;
 import static com.android.server.vcn.VcnTestUtils.setupIpSecManager;
 
 import static org.junit.Assert.assertEquals;
@@ -44,7 +45,6 @@
 import android.net.LinkAddress;
 import android.net.LinkProperties;
 import android.net.Network;
-import android.net.NetworkAgent;
 import android.net.NetworkCapabilities;
 import android.net.ipsec.ike.ChildSessionCallback;
 import android.net.ipsec.ike.IkeSessionCallback;
@@ -278,8 +278,8 @@
 
     protected void verifySafeModeTimeoutNotifiesCallbackAndUnregistersNetworkAgent(
             @NonNull State expectedState) {
-        // Set a NetworkAgent, and expect it to be unregistered and cleared
-        final NetworkAgent mockNetworkAgent = mock(NetworkAgent.class);
+        // Set a VcnNetworkAgent, and expect it to be unregistered and cleared
+        final VcnNetworkAgent mockNetworkAgent = mock(VcnNetworkAgent.class);
         mGatewayConnection.setNetworkAgent(mockNetworkAgent);
 
         // SafeMode timer starts when VcnGatewayConnection exits DisconnectedState (the initial