commit a8a74f19884bffc4b1dd029b2eefb3cae4467db6
author Aaron Huang <huangaaron@google.com> Tue Dec 22 23:29:04 2020 +0800
committer Aaron Huang <huangaaron@google.com> Thu Feb 04 19:14:57 2021 +0800
tree b46a13bfdd07c9892504d060785c976f3bf8807a
parent 068b629dfefa1fefcafd543d7900b963f2b720d5

Update VpnTest for new restricted API

This patch addressses the mocked UserManager for the context
of the user which is created by createContextAsUser.

Bug: 171529940
Bug: 175883995
Test: FrameworksNetTests
Change-Id: Ie72d08c765ed7fb7083c0c4b7ee02a43c23d1707

tests/net/java/com/android/server/connectivity/VpnTest.java

diff --git a/tests/net/java/com/android/server/connectivity/VpnTest.java b/tests/net/java/com/android/server/connectivity/VpnTest.java
index f478282..eff439d 100644
--- a/tests/net/java/com/android/server/connectivity/VpnTest.java
+++ b/tests/net/java/com/android/server/connectivity/VpnTest.java
@@ -253,12 +253,14 @@
 
     @Test
     public void testRestrictedProfilesAreAddedToVpn() {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         setMockedUsers(primaryUser, secondaryUser, restrictedProfileA, restrictedProfileB);
 
         final Vpn vpn = createVpn(primaryUser.id);
-        final Set<UidRange> ranges = vpn.createUserAndRestrictedProfilesRanges(primaryUser.id,
-                null, null);
+
+        // Assume the user can have restricted profiles.
+        doReturn(true).when(mUserManager).canHaveRestrictedProfile();
+        final Set<UidRange> ranges =
+                vpn.createUserAndRestrictedProfilesRanges(primaryUser.id, null, null);
 
         assertEquals(new ArraySet<>(Arrays.asList(new UidRange[] {
                 PRI_USER_RANGE, UidRange.createForUser(restrictedProfileA.id)
@@ -267,7 +269,6 @@
 
     @Test
     public void testManagedProfilesAreNotAddedToVpn() {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         setMockedUsers(primaryUser, managedProfileA);
 
         final Vpn vpn = createVpn(primaryUser.id);
@@ -290,7 +291,6 @@
 
     @Test
     public void testUidAllowAndDenylist() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRange user = PRI_USER_RANGE;
         final String[] packages = {PKGS[0], PKGS[1], PKGS[2]};
@@ -316,7 +316,6 @@
 
     @Test
     public void testGetAlwaysAndOnGetLockDown() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
 
         // Default state.
@@ -341,7 +340,6 @@
 
     @Test
     public void testLockdownChangingPackage() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRange user = PRI_USER_RANGE;
 
@@ -369,7 +367,6 @@
 
     @Test
     public void testLockdownAllowlist() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRange user = PRI_USER_RANGE;
 
@@ -444,7 +441,6 @@
 
     @Test
     public void testLockdownRuleRepeatability() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRangeParcel[] primaryUserRangeParcel = new UidRangeParcel[] {
                 new UidRangeParcel(PRI_USER_RANGE.start, PRI_USER_RANGE.stop)};
@@ -477,7 +473,6 @@
 
     @Test
     public void testLockdownRuleReversibility() throws Exception {
-        if (true) return; // TODO(b/175883995): Test disabled until updated for new UserManager API.
         final Vpn vpn = createVpn(primaryUser.id);
         final UidRangeParcel[] entireUser = {
             new UidRangeParcel(PRI_USER_RANGE.start, PRI_USER_RANGE.stop)
@@ -1144,6 +1139,10 @@
         doReturn(UserHandle.of(userId)).when(asUserContext).getUser();
         when(mContext.createContextAsUser(eq(UserHandle.of(userId)), anyInt()))
                 .thenReturn(asUserContext);
+        when(asUserContext.getSystemServiceName(UserManager.class))
+                .thenReturn(Context.USER_SERVICE);
+        when(asUserContext.getSystemService(UserManager.class))
+                .thenReturn(mUserManager);
         final TestLooper testLooper = new TestLooper();
         final Vpn vpn = new Vpn(testLooper.getLooper(), mContext, new TestDeps(), mNetService,
                 mNetd, userId, mKeyStore, mSystemServices, mIkev2SessionCreator);
@@ -1179,11 +1178,6 @@
             final int id = (int) invocation.getArguments()[0];
             return userMap.get(id);
         }).when(mUserManager).getUserInfo(anyInt());
-
-        doAnswer(invocation -> {
-            final int id = (int) invocation.getArguments()[0];
-            return (userMap.get(id).flags & UserInfo.FLAG_ADMIN) != 0;
-        }).when(mUserManager).canHaveRestrictedProfile();
     }
 
     /**