Add a test for new UID checking code.

Test: manual, this is a new test and it passes
Change-Id: I69a39039647ec05f970700e36c5f1e1e6e110f4a
diff --git a/tests/net/java/com/android/server/ConnectivityServiceTest.java b/tests/net/java/com/android/server/ConnectivityServiceTest.java
index 70cacb3..4733deb 100644
--- a/tests/net/java/com/android/server/ConnectivityServiceTest.java
+++ b/tests/net/java/com/android/server/ConnectivityServiceTest.java
@@ -44,6 +44,7 @@
 import static android.net.NetworkCapabilities.NET_CAPABILITY_XCAP;
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
 import static android.net.NetworkCapabilities.TRANSPORT_ETHERNET;
+import static android.net.NetworkCapabilities.TRANSPORT_VPN;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI_AWARE;
 
@@ -101,6 +102,7 @@
 import android.net.NetworkUtils;
 import android.net.RouteInfo;
 import android.net.StringNetworkSpecifier;
+import android.net.UidRange;
 import android.net.metrics.IpConnectivityLog;
 import android.net.util.MultinetworkPolicyTracker;
 import android.os.ConditionVariable;
@@ -360,7 +362,7 @@
 
         MockNetworkAgent(int transport, LinkProperties linkProperties) {
             final int type = transportToLegacyType(transport);
-            final String typeName = ConnectivityManager.getNetworkTypeName(type);
+            final String typeName = ConnectivityManager.getNetworkTypeName(transport);
             mNetworkInfo = new NetworkInfo(type, 0, typeName, "Mock");
             mNetworkCapabilities = new NetworkCapabilities();
             mNetworkCapabilities.addTransportType(transport);
@@ -377,6 +379,9 @@
                 case TRANSPORT_WIFI_AWARE:
                     mScore = 20;
                     break;
+                case TRANSPORT_VPN:
+                    mScore = 0;
+                    break;
                 default:
                     throw new UnsupportedOperationException("unimplemented network type");
             }
@@ -438,6 +443,11 @@
             mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
         }
 
+        public void setUids(Set<UidRange> uids) {
+            mNetworkCapabilities.setUids(uids);
+            mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
+        }
+
         public void setSignalStrength(int signalStrength) {
             mNetworkCapabilities.setSignalStrength(signalStrength);
             mNetworkAgent.sendNetworkCapabilities(mNetworkCapabilities);
@@ -1463,6 +1473,11 @@
             return nc;
         }
 
+        void expectCapabilitiesLike(Predicate<NetworkCapabilities> fn, MockNetworkAgent agent) {
+            CallbackInfo cbi = expectCallback(CallbackState.NETWORK_CAPABILITIES, agent);
+            assertTrue(fn.test((NetworkCapabilities) cbi.arg));
+        }
+
         void assertNoCallback() {
             waitForIdle();
             CallbackInfo c = mCallbacks.peek();
@@ -3625,4 +3640,76 @@
             return;
         }
     }
+
+    @Test
+    public void testVpnNetworkActive() {
+        final int uid = Process.myUid();
+
+        final TestNetworkCallback genericNetworkCallback = new TestNetworkCallback();
+        final TestNetworkCallback wifiNetworkCallback = new TestNetworkCallback();
+        final TestNetworkCallback vpnNetworkCallback = new TestNetworkCallback();
+        final NetworkRequest genericRequest = new NetworkRequest.Builder().build();
+        final NetworkRequest wifiRequest = new NetworkRequest.Builder()
+                .addTransportType(TRANSPORT_WIFI).build();
+        final NetworkRequest vpnNetworkRequest = new NetworkRequest.Builder()
+                .addTransportType(TRANSPORT_VPN).build();
+        mCm.registerNetworkCallback(genericRequest, genericNetworkCallback);
+        mCm.registerNetworkCallback(wifiRequest, wifiNetworkCallback);
+        mCm.registerNetworkCallback(vpnNetworkRequest, vpnNetworkCallback);
+
+        mWiFiNetworkAgent = new MockNetworkAgent(TRANSPORT_WIFI);
+        mWiFiNetworkAgent.connect(false);
+
+        genericNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        wifiNetworkCallback.expectAvailableCallbacksUnvalidated(mWiFiNetworkAgent);
+        vpnNetworkCallback.assertNoCallback();
+
+        // TODO : check callbacks agree with the return value of mCm.getActiveNetwork().
+        // Right now this is not possible because establish() is not adequately instrumented
+        // in this test.
+
+        final MockNetworkAgent vpnNetworkAgent = new MockNetworkAgent(TRANSPORT_VPN);
+        final ArraySet<UidRange> ranges = new ArraySet<>();
+        ranges.add(new UidRange(uid, uid));
+        vpnNetworkAgent.setUids(ranges);
+        vpnNetworkAgent.connect(false);
+
+        genericNetworkCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
+        wifiNetworkCallback.assertNoCallback();
+        vpnNetworkCallback.expectAvailableCallbacksUnvalidated(vpnNetworkAgent);
+
+        genericNetworkCallback.expectCallback(CallbackState.NETWORK_CAPABILITIES, vpnNetworkAgent);
+        vpnNetworkCallback.expectCapabilitiesLike(
+                nc -> nc.appliesToUid(uid) && !nc.appliesToUid(uid + 1), vpnNetworkAgent);
+
+        ranges.clear();
+        vpnNetworkAgent.setUids(ranges);
+
+        genericNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        wifiNetworkCallback.assertNoCallback();
+        vpnNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+
+        ranges.add(new UidRange(uid, uid));
+        vpnNetworkAgent.setUids(ranges);
+
+        genericNetworkCallback.expectAvailableCallbacksValidated(vpnNetworkAgent);
+        wifiNetworkCallback.assertNoCallback();
+        vpnNetworkCallback.expectAvailableCallbacksValidated(vpnNetworkAgent);
+
+        mWiFiNetworkAgent.disconnect();
+
+        genericNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        wifiNetworkCallback.expectCallback(CallbackState.LOST, mWiFiNetworkAgent);
+        vpnNetworkCallback.assertNoCallback();
+
+        vpnNetworkAgent.disconnect();
+
+        genericNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+        wifiNetworkCallback.assertNoCallback();
+        vpnNetworkCallback.expectCallback(CallbackState.LOST, vpnNetworkAgent);
+
+        mCm.unregisterNetworkCallback(genericNetworkCallback);
+        mCm.unregisterNetworkCallback(wifiNetworkCallback);
+        mCm.unregisterNetworkCallback(vpnNetworkCallback);
+    }
 }