1 /*
2  * Copyright (C) 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package android.virtualdevice.cts.common;
18 
19 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
20 
21 import static com.google.common.truth.Truth.assertThat;
22 
23 import static org.junit.Assume.assumeTrue;
24 import static org.mockito.ArgumentMatchers.any;
25 import static org.mockito.ArgumentMatchers.eq;
26 import static org.mockito.Mockito.mock;
27 import static org.mockito.Mockito.reset;
28 import static org.mockito.Mockito.timeout;
29 import static org.mockito.Mockito.verify;
30 
31 import android.app.role.RoleManager;
32 import android.companion.AssociationInfo;
33 import android.companion.AssociationRequest;
34 import android.companion.CompanionDeviceManager;
35 import android.content.Context;
36 import android.content.pm.PackageManager;
37 import android.os.Process;
38 import android.util.Log;
39 
40 import com.android.compatibility.common.util.FeatureUtil;
41 import com.android.compatibility.common.util.SystemUtil;
42 import com.android.modules.utils.build.SdkLevel;
43 
44 import org.junit.rules.ExternalResource;
45 import org.mockito.Mock;
46 import org.mockito.MockitoAnnotations;
47 
48 import java.util.List;
49 import java.util.Locale;
50 import java.util.concurrent.Executor;
51 import java.util.function.Consumer;
52 
53 /**
54  * A test rule that creates a {@link CompanionDeviceManager} association with the instrumented
55  * package for the duration of the test.
56  */
57 public class FakeAssociationRule extends ExternalResource {
58     private static final String TAG = "FakeAssociationRule";
59 
60     private static final String FAKE_ASSOCIATION_ADDRESS_FORMAT = "00:00:00:00:00:%02d";
61 
62     private static final int TIMEOUT_MS = 10000;
63 
64     private final Context mContext = getInstrumentation().getTargetContext();
65 
66     private final Executor mCallbackExecutor = Runnable::run;
67     private final RoleManager mRoleManager = mContext.getSystemService(RoleManager.class);
68 
69     private final String mDeviceProfile;
70 
71     @Mock
72     private CompanionDeviceManager.OnAssociationsChangedListener mOnAssociationsChangedListener;
73 
74     private int mNextDeviceId = 0;
75 
76     private AssociationInfo mAssociationInfo;
77     private CompanionDeviceManager mCompanionDeviceManager;
78 
FakeAssociationRule()79     public FakeAssociationRule() {
80         this(AssociationRequest.DEVICE_PROFILE_APP_STREAMING);
81     }
82 
FakeAssociationRule(String deviceProfile)83     public FakeAssociationRule(String deviceProfile) {
84         mDeviceProfile = deviceProfile;
85         mCompanionDeviceManager = mContext.getSystemService(CompanionDeviceManager.class);
86     }
87 
createManagedAssociation()88     public AssociationInfo createManagedAssociation() {
89         String deviceAddress = String.format(Locale.getDefault(Locale.Category.FORMAT),
90                 FAKE_ASSOCIATION_ADDRESS_FORMAT, ++mNextDeviceId);
91         if (mNextDeviceId > 99) {
92             throw new IllegalArgumentException("At most 99 associations supported");
93         }
94         if (mNextDeviceId > 1 && !SdkLevel.isAtLeastT()) {
95             throw new IllegalArgumentException("Multiple associations require API level 33");
96         }
97 
98         Log.d(TAG, "Associations before shell cmd: "
99                 + mCompanionDeviceManager.getMyAssociations().size());
100         reset(mOnAssociationsChangedListener);
101         SystemUtil.runShellCommandOrThrow(String.format(Locale.getDefault(Locale.Category.FORMAT),
102                 "cmd companiondevice associate %d %s %s %s",
103                 getInstrumentation().getContext().getUserId(),
104                 mContext.getPackageName(),
105                 deviceAddress,
106                 mDeviceProfile));
107         verify(mOnAssociationsChangedListener, timeout(TIMEOUT_MS)
108                 .description(TAG
109                         + ": Association changed listener did not call back. Total associations: "
110                         + mCompanionDeviceManager.getMyAssociations().size()))
111                 .onAssociationsChanged(any());
112         List<AssociationInfo> associations = mCompanionDeviceManager.getMyAssociations();
113 
114         if (SdkLevel.isAtLeastT()) {
115             final AssociationInfo associationInfo = associations.stream()
116                     .filter(a -> deviceAddress.equals(a.getDeviceMacAddress().toString()))
117                     .findAny().orElse(null);
118             assertThat(associationInfo).isNotNull();
119             return associationInfo;
120         } else {
121             assertThat(associations).hasSize(1);
122             return associations.get(0);
123         }
124     }
125 
126     @Override
before()127     protected void before() throws Throwable {
128         super.before();
129         MockitoAnnotations.initMocks(this);
130         assumeTrue(FeatureUtil.hasSystemFeature(PackageManager.FEATURE_COMPANION_DEVICE_SETUP));
131 
132         Consumer<Boolean> callback = mock(Consumer.class);
133         SystemUtil.runWithShellPermissionIdentity(() -> {
134             mCompanionDeviceManager.addOnAssociationsChangedListener(
135                     mCallbackExecutor, mOnAssociationsChangedListener);
136             mRoleManager.setBypassingRoleQualification(true);
137             mRoleManager.addRoleHolderAsUser(
138                     mDeviceProfile, mContext.getPackageName(),
139                     RoleManager.MANAGE_HOLDERS_FLAG_DONT_KILL_APP, Process.myUserHandle(),
140                     mCallbackExecutor, callback);
141             verify(callback, timeout(TIMEOUT_MS)).accept(eq(true));
142         });
143 
144         clearExistingAssociations();
145         mAssociationInfo = createManagedAssociation();
146     }
147 
148     @Override
after()149     protected void after() {
150         super.after();
151         clearExistingAssociations();
152 
153         Consumer<Boolean> callback = mock(Consumer.class);
154         SystemUtil.runWithShellPermissionIdentity(() -> {
155             mRoleManager.removeRoleHolderAsUser(
156                     mDeviceProfile, mContext.getPackageName(),
157                     RoleManager.MANAGE_HOLDERS_FLAG_DONT_KILL_APP, Process.myUserHandle(),
158                     mCallbackExecutor, callback);
159             verify(callback, timeout(TIMEOUT_MS)).accept(eq(true));
160             mRoleManager.setBypassingRoleQualification(false);
161             mCompanionDeviceManager.removeOnAssociationsChangedListener(
162                     mOnAssociationsChangedListener);
163         });
164     }
165 
clearExistingAssociations()166     private void clearExistingAssociations() {
167         List<AssociationInfo> associations = mCompanionDeviceManager.getMyAssociations();
168         for (AssociationInfo association : associations) {
169             disassociate(association.getId());
170         }
171         assertThat(mCompanionDeviceManager.getMyAssociations()).isEmpty();
172         mAssociationInfo = null;
173     }
174 
getAssociationInfo()175     public AssociationInfo getAssociationInfo() {
176         return mAssociationInfo;
177     }
178 
disassociate()179     public void disassociate() {
180         clearExistingAssociations();
181     }
182 
disassociate(int associationId)183     private void disassociate(int associationId) {
184         reset(mOnAssociationsChangedListener);
185         mCompanionDeviceManager.disassociate(associationId);
186         verify(mOnAssociationsChangedListener, timeout(TIMEOUT_MS).atLeastOnce())
187             .onAssociationsChanged(any());
188     }
189 }
190