Merge "VCN: Allow restricting TRANSPORT_TEST for CTS"
diff --git a/core/java/android/net/vcn/VcnConfig.java b/core/java/android/net/vcn/VcnConfig.java
index 6f9c9dd..a27e923 100644
--- a/core/java/android/net/vcn/VcnConfig.java
+++ b/core/java/android/net/vcn/VcnConfig.java
@@ -16,6 +16,7 @@
 package android.net.vcn;
 
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
+import static android.net.NetworkCapabilities.TRANSPORT_TEST;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 
 import static com.android.internal.annotations.VisibleForTesting.Visibility;
@@ -75,6 +76,7 @@
     static {
         ALLOWED_TRANSPORTS.add(TRANSPORT_WIFI);
         ALLOWED_TRANSPORTS.add(TRANSPORT_CELLULAR);
+        ALLOWED_TRANSPORTS.add(TRANSPORT_TEST);
     }
 
     private static final String PACKAGE_NAME_KEY = "mPackageName";
@@ -155,6 +157,11 @@
                                 + transport
                                 + " which might be from a new version of VcnConfig");
             }
+
+            if (transport == TRANSPORT_TEST && !mIsTestModeProfile) {
+                throw new IllegalArgumentException(
+                        "Found TRANSPORT_TEST in a non-test-mode profile");
+            }
         }
     }
 
diff --git a/services/core/java/com/android/server/VcnManagementService.java b/services/core/java/com/android/server/VcnManagementService.java
index 78d4708..e8c85ce 100644
--- a/services/core/java/com/android/server/VcnManagementService.java
+++ b/services/core/java/com/android/server/VcnManagementService.java
@@ -1074,9 +1074,10 @@
                             subGrp, mLastSnapshot, mConfigs.get(subGrp));
                     for (int restrictedTransport : restrictedTransports) {
                         if (ncCopy.hasTransport(restrictedTransport)) {
-                            if (restrictedTransport == TRANSPORT_CELLULAR) {
-                                // Only make a cell network as restricted when the VCN is in
-                                // active mode.
+                            if (restrictedTransport == TRANSPORT_CELLULAR
+                                    || restrictedTransport == TRANSPORT_TEST) {
+                                // For cell or test network, only mark it as restricted when
+                                // the VCN is in active mode.
                                 isRestricted |= (vcn.getStatus() == VCN_STATUS_CODE_ACTIVE);
                             } else {
                                 isRestricted = true;
diff --git a/tests/vcn/java/android/net/vcn/VcnConfigTest.java b/tests/vcn/java/android/net/vcn/VcnConfigTest.java
index b313c9f..73a0a61 100644
--- a/tests/vcn/java/android/net/vcn/VcnConfigTest.java
+++ b/tests/vcn/java/android/net/vcn/VcnConfigTest.java
@@ -17,6 +17,7 @@
 package android.net.vcn;
 
 import static android.net.NetworkCapabilities.TRANSPORT_CELLULAR;
+import static android.net.NetworkCapabilities.TRANSPORT_TEST;
 import static android.net.NetworkCapabilities.TRANSPORT_WIFI;
 
 import static org.junit.Assert.assertEquals;
@@ -160,6 +161,37 @@
         assertNotEquals(config, configNotEqual);
     }
 
+    private VcnConfig buildConfigRestrictTransportTest(boolean isTestMode) throws Exception {
+        VcnConfig.Builder builder =
+                new VcnConfig.Builder(mContext)
+                        .setRestrictedUnderlyingNetworkTransports(Set.of(TRANSPORT_TEST));
+        if (isTestMode) {
+            builder.setIsTestModeProfile();
+        }
+
+        for (VcnGatewayConnectionConfig gatewayConnectionConfig : GATEWAY_CONNECTION_CONFIGS) {
+            builder.addGatewayConnectionConfig(gatewayConnectionConfig);
+        }
+
+        return builder.build();
+    }
+
+    @Test
+    public void testRestrictTransportTestInTestModeProfile() throws Exception {
+        final VcnConfig config = buildConfigRestrictTransportTest(true /*  isTestMode */);
+        assertEquals(Set.of(TRANSPORT_TEST), config.getRestrictedUnderlyingNetworkTransports());
+    }
+
+    @Test
+    public void testRestrictTransportTestInNonTestModeProfile() throws Exception {
+        try {
+            buildConfigRestrictTransportTest(false /*  isTestMode */);
+            fail("Expected exception because the config is not a test mode profile");
+        } catch (Exception expected) {
+
+        }
+    }
+
     @Test
     public void testParceling() {
         final VcnConfig config = buildTestConfig(mContext);