1 /*
2  * Copyright (C) 2022 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 package com.android.server.appsearch;
17 
18 import static android.Manifest.permission.READ_GLOBAL_APP_SEARCH_DATA;
19 import static android.app.appsearch.AppSearchResult.RESULT_DENIED;
20 import static android.app.appsearch.AppSearchResult.RESULT_RATE_LIMITED;
21 import static android.system.OsConstants.O_RDONLY;
22 import static android.system.OsConstants.O_WRONLY;
23 
24 import static com.android.internal.util.ConcurrentUtils.DIRECT_EXECUTOR;
25 import static com.android.server.appsearch.FrameworkServiceAppSearchConfig.KEY_DENYLIST;
26 import static com.android.server.appsearch.FrameworkServiceAppSearchConfig.KEY_RATE_LIMIT_API_COSTS;
27 import static com.android.server.appsearch.FrameworkServiceAppSearchConfig.KEY_RATE_LIMIT_ENABLED;
28 import static com.android.server.appsearch.FrameworkServiceAppSearchConfig.KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE;
29 import static com.android.server.appsearch.FrameworkServiceAppSearchConfig.KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY;
30 
31 import static com.google.common.truth.Truth.assertThat;
32 
33 import static org.mockito.ArgumentMatchers.any;
34 import static org.mockito.ArgumentMatchers.anyBoolean;
35 import static org.mockito.ArgumentMatchers.anyInt;
36 import static org.mockito.ArgumentMatchers.anyString;
37 import static org.mockito.ArgumentMatchers.eq;
38 import static org.mockito.Mockito.clearInvocations;
39 import static org.mockito.Mockito.doReturn;
40 import static org.mockito.Mockito.mock;
41 import static org.mockito.Mockito.never;
42 import static org.mockito.Mockito.spy;
43 import static org.mockito.Mockito.timeout;
44 import static org.mockito.Mockito.verify;
45 
46 import android.Manifest;
47 import android.annotation.NonNull;
48 import android.annotation.Nullable;
49 import android.app.UiAutomation;
50 import android.app.admin.DevicePolicyManager;
51 import android.app.appsearch.AppSearchBatchResult;
52 import android.app.appsearch.AppSearchEnvironmentFactory;
53 import android.app.appsearch.AppSearchResult;
54 import android.app.appsearch.AppSearchSchema;
55 import android.app.appsearch.AppSearchSchema.LongPropertyConfig;
56 import android.app.appsearch.AppSearchSchema.PropertyConfig;
57 import android.app.appsearch.AppSearchSchema.StringPropertyConfig;
58 import android.app.appsearch.FrameworkAppSearchEnvironment;
59 import android.app.appsearch.GenericDocument;
60 import android.app.appsearch.GetByDocumentIdRequest;
61 import android.app.appsearch.GetSchemaResponse;
62 import android.app.appsearch.InternalSetSchemaResponse;
63 import android.app.appsearch.RemoveByDocumentIdRequest;
64 import android.app.appsearch.ReportUsageRequest;
65 import android.app.appsearch.SearchResultPage;
66 import android.app.appsearch.SearchSpec;
67 import android.app.appsearch.SearchSuggestionSpec;
68 import android.app.appsearch.aidl.AppSearchAttributionSource;
69 import android.app.appsearch.aidl.AppSearchBatchResultParcel;
70 import android.app.appsearch.aidl.AppSearchResultParcel;
71 import android.app.appsearch.aidl.DocumentsParcel;
72 import android.app.appsearch.aidl.ExecuteAppFunctionAidlRequest;
73 import android.app.appsearch.aidl.GetDocumentsAidlRequest;
74 import android.app.appsearch.aidl.GetNamespacesAidlRequest;
75 import android.app.appsearch.aidl.GetNextPageAidlRequest;
76 import android.app.appsearch.aidl.GetSchemaAidlRequest;
77 import android.app.appsearch.aidl.GetStorageInfoAidlRequest;
78 import android.app.appsearch.aidl.GlobalSearchAidlRequest;
79 import android.app.appsearch.aidl.IAppFunctionService;
80 import android.app.appsearch.aidl.IAppSearchBatchResultCallback;
81 import android.app.appsearch.aidl.IAppSearchManager;
82 import android.app.appsearch.aidl.IAppSearchObserverProxy;
83 import android.app.appsearch.aidl.IAppSearchResultCallback;
84 import android.app.appsearch.aidl.InitializeAidlRequest;
85 import android.app.appsearch.aidl.InvalidateNextPageTokenAidlRequest;
86 import android.app.appsearch.aidl.PersistToDiskAidlRequest;
87 import android.app.appsearch.aidl.PutDocumentsAidlRequest;
88 import android.app.appsearch.aidl.PutDocumentsFromFileAidlRequest;
89 import android.app.appsearch.aidl.RegisterObserverCallbackAidlRequest;
90 import android.app.appsearch.aidl.RemoveByDocumentIdAidlRequest;
91 import android.app.appsearch.aidl.RemoveByQueryAidlRequest;
92 import android.app.appsearch.aidl.ReportUsageAidlRequest;
93 import android.app.appsearch.aidl.SearchAidlRequest;
94 import android.app.appsearch.aidl.SearchSuggestionAidlRequest;
95 import android.app.appsearch.aidl.SetSchemaAidlRequest;
96 import android.app.appsearch.aidl.UnregisterObserverCallbackAidlRequest;
97 import android.app.appsearch.aidl.WriteSearchResultsToFileAidlRequest;
98 import android.app.appsearch.functions.AppFunctionManager;
99 import android.app.appsearch.functions.ExecuteAppFunctionRequest;
100 import android.app.appsearch.functions.ExecuteAppFunctionResponse;
101 import android.app.appsearch.functions.ServiceCallHelper;
102 import android.app.appsearch.observer.ObserverSpec;
103 import android.app.appsearch.safeparcel.GenericDocumentParcel;
104 import android.app.appsearch.stats.SchemaMigrationStats;
105 import android.app.role.RoleManager;
106 import android.content.AttributionSource;
107 import android.content.BroadcastReceiver;
108 import android.content.Context;
109 import android.content.ContextWrapper;
110 import android.content.Intent;
111 import android.content.IntentFilter;
112 import android.content.pm.PackageManager;
113 import android.content.pm.ResolveInfo;
114 import android.content.pm.ServiceInfo;
115 import android.os.Handler;
116 import android.os.ParcelFileDescriptor;
117 import android.os.RemoteException;
118 import android.os.ServiceManager;
119 import android.os.SystemClock;
120 import android.os.UserHandle;
121 import android.os.UserManager;
122 import android.provider.DeviceConfig;
123 
124 import androidx.test.core.app.ApplicationProvider;
125 import androidx.test.platform.app.InstrumentationRegistry;
126 
127 import com.android.dx.mockito.inline.extended.ExtendedMockito;
128 import com.android.dx.mockito.inline.extended.StaticMockitoSessionBuilder;
129 import com.android.modules.utils.testing.ExtendedMockitoRule;
130 import com.android.modules.utils.testing.StaticMockFixture;
131 import com.android.modules.utils.testing.TestableDeviceConfig;
132 import com.android.server.LocalManagerRegistry;
133 import com.android.server.appsearch.external.localstorage.stats.CallStats;
134 import com.android.server.appsearch.external.localstorage.stats.SearchIntentStats;
135 import com.android.server.appsearch.external.localstorage.stats.SearchSessionStats;
136 import com.android.server.appsearch.external.localstorage.stats.SearchStats;
137 import com.android.server.appsearch.external.localstorage.stats.SetSchemaStats;
138 import com.android.server.appsearch.external.localstorage.usagereporting.ClickActionGenericDocument;
139 import com.android.server.appsearch.external.localstorage.usagereporting.SearchActionGenericDocument;
140 import com.android.server.usage.StorageStatsManagerLocal;
141 
142 import com.google.common.util.concurrent.SettableFuture;
143 
144 import libcore.io.IoBridge;
145 
146 import org.junit.After;
147 import org.junit.Before;
148 import org.junit.Rule;
149 import org.junit.Test;
150 import org.junit.rules.TemporaryFolder;
151 import org.mockito.ArgumentCaptor;
152 
153 import java.io.File;
154 import java.io.FileDescriptor;
155 import java.util.Arrays;
156 import java.util.Collections;
157 import java.util.List;
158 import java.util.Objects;
159 import java.util.concurrent.ExecutionException;
160 import java.util.function.Consumer;
161 
162 public class AppSearchManagerServiceTest {
163     private static final String DATABASE_NAME = "databaseName";
164     private static final String NAMESPACE = "namespace";
165     private static final String ID = "ID";
166     private static final SearchSpec EMPTY_SEARCH_SPEC = new SearchSpec.Builder().build();
167     // Mostly guarantees the logged estimated binder latency is positive and doesn't overflow
168     private static final long BINDER_CALL_START_TIME = SystemClock.elapsedRealtime() - 1;
169     private static final String FOO_PACKAGE_NAME = "foo";
170 
171     private final MockServiceManager mMockServiceManager = new MockServiceManager();
172     private final RoleManager mRoleManager = mock(RoleManager.class);
173     private final DevicePolicyManager mDevicePolicyManager = mock(DevicePolicyManager.class);
174 
175     @Rule
176     public ExtendedMockitoRule mExtendedMockitoRule = new ExtendedMockitoRule.Builder()
177             .addStaticMockFixtures(() -> mMockServiceManager, TestableDeviceConfig::new)
178             .build();
179 
180     @Rule
181     public TemporaryFolder mTemporaryFolder = new TemporaryFolder();
182 
183     private Context mContext;
184     private AppSearchManagerService mAppSearchManagerService;
185     private UserHandle mUserHandle;
186     private UiAutomation mUiAutomation;
187     private IAppSearchManager.Stub mAppSearchManagerServiceStub;
188     private AppSearchUserInstance mUserInstance;
189     private InternalAppSearchLogger mLogger;
190     private TestableServiceCallHelper mServiceCallHelper;
191 
192     private int mCallingPid;
193 
194     @Before
setUp()195     public void setUp() throws Exception {
196         Context context = ApplicationProvider.getApplicationContext();
197         mUserHandle = context.getUser();
198         mUiAutomation = InstrumentationRegistry.getInstrumentation().getUiAutomation();
199         mContext = new ContextWrapper(context) {
200             // Mock-able package manager for testing
201             final PackageManager mPackageManager = spy(context.getPackageManager());
202             final UserManager mUserManager = spy(context.getSystemService(UserManager.class));
203 
204             @Override
205             public Intent registerReceiverForAllUsers(@Nullable BroadcastReceiver receiver,
206                     @NonNull IntentFilter filter, @Nullable String broadcastPermission,
207                     @Nullable Handler scheduler) {
208                 // Do nothing
209                 return null;
210             }
211 
212             @Override
213             public Context createContextAsUser(UserHandle user, int flags) {
214                 return new ContextWrapper(super.createContextAsUser(user, flags)) {
215                     @Override
216                     public PackageManager getPackageManager() {
217                         return mPackageManager;
218                     }
219                 };
220             }
221 
222             @Override
223             public PackageManager getPackageManager() {
224                 return mPackageManager;
225             }
226 
227             @Nullable
228             @Override
229             public Object getSystemService(String name) {
230                 if (Context.ROLE_SERVICE.equals(name)) {
231                     return mRoleManager;
232                 }
233                 if (Context.DEVICE_POLICY_SERVICE.equals(name)) {
234                     return mDevicePolicyManager;
235                 }
236                 if (Context.USER_SERVICE.equals(name)) {
237                     return mUserManager;
238                 }
239                 return super.getSystemService(name);
240             }
241         };
242 
243         // Set a test environment that provides a temporary folder for AppSearch
244         File mAppSearchDir = mTemporaryFolder.newFolder();
245         AppSearchEnvironmentFactory.setEnvironmentInstanceForTest(
246                 new FrameworkAppSearchEnvironment() {
247                     @Override
248                     public File getAppSearchDir(@NonNull Context unused,
249                             @NonNull UserHandle userHandle) {
250                         return mAppSearchDir;
251                     }
252                 });
253 
254         setUpEnvironmentForAppFunction();
255         mServiceCallHelper = new TestableServiceCallHelper();
256 
257         // In AppSearchManagerService, ServiceAppSearchConfig is a singleton. During tearDown for
258         // TestableDeviceConfig, the propertyChangedListeners are removed. Therefore we have to set
259         // a fresh config with listeners in setUp in order to set new properties.
260         ServiceAppSearchConfig appSearchConfig =
261                 FrameworkServiceAppSearchConfig.create(DIRECT_EXECUTOR);
262         AppSearchComponentFactory.setConfigInstanceForTest(appSearchConfig);
263 
264         // Create the user instance and add a spy to its logger to verify logging
265         // Note, SimpleTestLogger does not suffice for our tests since CallStats logging in
266         // AppSearchManagerService occurs in a separate thread. With a spy, we can verify with a
267         // timeout to catch asynchronous calls.
268         mUserInstance = AppSearchUserInstanceManager.getInstance().getOrCreateUserInstance(mContext,
269                 mUserHandle, appSearchConfig);
270         mLogger = spy(mUserInstance.getLogger());
271         mUserInstance.setLoggerForTest(mLogger);
272 
273         // Start the service
274         mAppSearchManagerService = new AppSearchManagerService(mContext,
275                 new AppSearchModule.Lifecycle(mContext), mServiceCallHelper);
276         mAppSearchManagerService.onStart();
277         mAppSearchManagerServiceStub = mMockServiceManager.mStubCaptor.getValue();
278         assertThat(mAppSearchManagerServiceStub).isNotNull();
279         mCallingPid = android.os.Process.myPid();
280     }
281 
282     @After
tearDown()283     public void tearDown() {
284         // The TemporaryFolder rule's teardown will delete the current test folder; by removing the
285         // current user instance, the next test will be able to create a new AppSearchImpl with
286         // a new test folder
287         AppSearchUserInstanceManager.getInstance().closeAndRemoveUserInstance(mUserHandle);
288     }
289 
290     @Test
testCallingPackageDoesntExistsInTargetUser()291     public void testCallingPackageDoesntExistsInTargetUser() throws Exception {
292         UserHandle testTargetUser = new UserHandle(1234);
293         mUiAutomation.adoptShellPermissionIdentity(Manifest.permission.INTERACT_ACROSS_USERS_FULL);
294         try {
295             // Try to initial a AppSearchSession for secondary user, but the calling package doesn't
296             // exist in there.
297             TestResultCallback callback = new TestResultCallback();
298             mAppSearchManagerServiceStub.initialize(
299                     new InitializeAidlRequest(
300                             AppSearchAttributionSource.createAttributionSource(mContext,
301                                     mCallingPid),
302                             testTargetUser, System.currentTimeMillis())
303                     , callback);
304             assertThat(callback.get().isSuccess()).isFalse();
305             assertThat(callback.get().getErrorMessage()).contains(
306                     "SecurityException: Package: " + mContext.getPackageName()
307                             + " haven't installed for user " + testTargetUser.getIdentifier());
308         } finally {
309             mUiAutomation.dropShellPermissionIdentity();
310         }
311     }
312 
313     @Test
testSetSchemaStatsLogging()314     public void testSetSchemaStatsLogging() throws Exception {
315         TestResultCallback callback = new TestResultCallback();
316         mAppSearchManagerServiceStub.setSchema(
317                 new SetSchemaAidlRequest(
318                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
319                         DATABASE_NAME,
320                 /* schemaBundles= */ Collections.emptyList(),
321                 /* visibilityBundles= */ Collections.emptyList(), /* forceOverride= */ false,
322                 /* schemaVersion= */ 0, mUserHandle, BINDER_CALL_START_TIME,
323                 SchemaMigrationStats.FIRST_CALL_GET_INCOMPATIBLE),
324                 callback);
325         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
326         verifyCallStats(mContext.getPackageName(), DATABASE_NAME, CallStats.CALL_TYPE_SET_SCHEMA);
327         // SetSchemaStats is also logged in SetSchema
328         ArgumentCaptor<SetSchemaStats> setSchemaStatsCaptor = ArgumentCaptor.forClass(
329                 SetSchemaStats.class);
330         verify(mLogger, timeout(1000).times(1)).logStats(setSchemaStatsCaptor.capture());
331         SetSchemaStats setSchemaStats = setSchemaStatsCaptor.getValue();
332         assertThat(setSchemaStats.getPackageName()).isEqualTo(mContext.getPackageName());
333         assertThat(setSchemaStats.getDatabase()).isEqualTo(DATABASE_NAME);
334         assertThat(setSchemaStats.getStatusCode()).isEqualTo(callback.get().getResultCode());
335         assertThat(setSchemaStats.getSchemaMigrationCallType()).isEqualTo(
336                 SchemaMigrationStats.FIRST_CALL_GET_INCOMPATIBLE);
337     }
338 
339     @Test
testLocalGetSchemaStatsLogging()340     public void testLocalGetSchemaStatsLogging() throws Exception {
341         TestResultCallback callback = new TestResultCallback();
342         mAppSearchManagerServiceStub.getSchema(
343                 new GetSchemaAidlRequest(
344                         AppSearchAttributionSource.createAttributionSource(mContext,
345                                 mCallingPid),
346                         mContext.getPackageName(), DATABASE_NAME, mUserHandle,
347                         BINDER_CALL_START_TIME, /* isForEnterprise= */ false),
348                 callback);
349         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
350         verifyCallStats(mContext.getPackageName(), DATABASE_NAME, CallStats.CALL_TYPE_GET_SCHEMA);
351     }
352 
353     @Test
testGlobalGetSchemaStatsLogging()354     public void testGlobalGetSchemaStatsLogging() throws Exception {
355         String otherPackageName = mContext.getPackageName() + "foo";
356         TestResultCallback callback = new TestResultCallback();
357         mAppSearchManagerServiceStub.getSchema(
358                 new GetSchemaAidlRequest(
359                         AppSearchAttributionSource.createAttributionSource(mContext,
360                                 mCallingPid),
361                         otherPackageName, DATABASE_NAME, mUserHandle, BINDER_CALL_START_TIME,
362                         /* isForEnterprise= */ false),
363                 callback);
364         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
365         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
366                 CallStats.CALL_TYPE_GLOBAL_GET_SCHEMA);
367     }
368 
369     @Test
testGetNamespacesStatsLogging()370     public void testGetNamespacesStatsLogging() throws Exception {
371         TestResultCallback callback = new TestResultCallback();
372         mAppSearchManagerServiceStub.getNamespaces(
373                 new GetNamespacesAidlRequest(
374                         AppSearchAttributionSource.createAttributionSource(mContext,
375                                 mCallingPid), DATABASE_NAME,
376                         mUserHandle, BINDER_CALL_START_TIME),
377                 callback);
378         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
379         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
380                 CallStats.CALL_TYPE_GET_NAMESPACES);
381     }
382 
383     @Test
testPutDocumentsStatsLogging()384     public void testPutDocumentsStatsLogging() throws Exception {
385         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
386         mAppSearchManagerServiceStub.putDocuments(
387                 new PutDocumentsAidlRequest(
388                         AppSearchAttributionSource.createAttributionSource(mContext,
389                                 mCallingPid), DATABASE_NAME,
390                         new DocumentsParcel(Collections.emptyList(), Collections.emptyList()),
391                         mUserHandle, BINDER_CALL_START_TIME), callback);
392         assertThat(callback.get()).isNull(); // null means there wasn't an error
393         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
394                 CallStats.CALL_TYPE_PUT_DOCUMENTS);
395         // putDocuments only logs PutDocumentStats indirectly so we don't verify it
396     }
397 
398     @Test
testPutDocumentsStatsLogging_takenActions()399     public void testPutDocumentsStatsLogging_takenActions() throws Exception {
400         // Set SearchAction and ClickAction schemas.
401         List<AppSearchSchema> schemas =
402                 Arrays.asList(
403                         new AppSearchSchema.Builder("builtin:SearchAction")
404                                 .addProperty(
405                                         new LongPropertyConfig.Builder("actionType")
406                                                 .setCardinality(PropertyConfig.CARDINALITY_OPTIONAL)
407                                                 .build())
408                                 .addProperty(
409                                         new StringPropertyConfig.Builder("query")
410                                                 .setCardinality(PropertyConfig.CARDINALITY_OPTIONAL)
411                                                 .build())
412                                 .addProperty(
413                                         new LongPropertyConfig.Builder("fetchedResultCount")
414                                                 .setCardinality(PropertyConfig.CARDINALITY_OPTIONAL)
415                                                 .build())
416                                 .build(),
417                         new AppSearchSchema.Builder("builtin:ClickAction")
418                                 .addProperty(
419                                         new LongPropertyConfig.Builder("actionType")
420                                                 .setCardinality(PropertyConfig.CARDINALITY_OPTIONAL)
421                                                 .build())
422                                 .addProperty(
423                                         new StringPropertyConfig.Builder("query")
424                                                 .setCardinality(PropertyConfig.CARDINALITY_OPTIONAL)
425                                                 .build())
426                                 .addProperty(
427                                         new LongPropertyConfig.Builder("resultRankInBlock")
428                                                 .setCardinality(PropertyConfig.CARDINALITY_OPTIONAL)
429                                                 .build())
430                                 .addProperty(
431                                         new LongPropertyConfig.Builder("resultRankGlobal")
432                                                 .setCardinality(PropertyConfig.CARDINALITY_OPTIONAL)
433                                                 .build())
434                                 .addProperty(
435                                         new LongPropertyConfig.Builder("timeStayOnResultMillis")
436                                                 .setCardinality(PropertyConfig.CARDINALITY_OPTIONAL)
437                                                 .build())
438                                 .build());
439         InternalSetSchemaResponse internalSetSchemaResponse =
440                 mUserInstance
441                         .getAppSearchImpl()
442                         .setSchema(
443                                 mContext.getPackageName(),
444                                 DATABASE_NAME,
445                                 schemas,
446                                 /* visibilityDocuments= */ Collections.emptyList(),
447                                 /* forceOverride= */ false,
448                                 /* version= */ 0,
449                                 /* setSchemaStatsBuilder= */ null);
450         assertThat(internalSetSchemaResponse.isSuccess()).isTrue();
451 
452         // Prepare search action and click action generic documents.
453         SearchActionGenericDocument searchAction1 =
454                 new SearchActionGenericDocument.Builder(
455                                 "namespace", "search1", "builtin:SearchAction")
456                         .setCreationTimestampMillis(1000)
457                         .setQuery("tes")
458                         .setFetchedResultCount(20)
459                         .build();
460         ClickActionGenericDocument clickAction1 =
461                 new ClickActionGenericDocument.Builder("namespace", "click1", "builtin:ClickAction")
462                         .setCreationTimestampMillis(2000)
463                         .setQuery("tes")
464                         .setResultRankInBlock(1)
465                         .setResultRankGlobal(2)
466                         .setTimeStayOnResultMillis(512)
467                         .build();
468         ClickActionGenericDocument clickAction2 =
469                 new ClickActionGenericDocument.Builder("namespace", "click2", "builtin:ClickAction")
470                         .setCreationTimestampMillis(3000)
471                         .setQuery("tes")
472                         .setResultRankInBlock(3)
473                         .setResultRankGlobal(6)
474                         .setTimeStayOnResultMillis(1024)
475                         .build();
476         SearchActionGenericDocument searchAction2 =
477                 new SearchActionGenericDocument.Builder(
478                                 "namespace", "search2", "builtin:SearchAction")
479                         .setCreationTimestampMillis(5000)
480                         .setQuery("test")
481                         .setFetchedResultCount(10)
482                         .build();
483         ClickActionGenericDocument clickAction3 =
484                 new ClickActionGenericDocument.Builder("namespace", "click3", "builtin:ClickAction")
485                         .setCreationTimestampMillis(6000)
486                         .setQuery("test")
487                         .setResultRankInBlock(2)
488                         .setResultRankGlobal(4)
489                         .setTimeStayOnResultMillis(512)
490                         .build();
491         List<GenericDocumentParcel> takenActionGenericDocumentParcels =
492                 Arrays.asList(
493                         GenericDocumentParcel.fromGenericDocument(searchAction1),
494                         GenericDocumentParcel.fromGenericDocument(clickAction1),
495                         GenericDocumentParcel.fromGenericDocument(clickAction2),
496                         GenericDocumentParcel.fromGenericDocument(searchAction2),
497                         GenericDocumentParcel.fromGenericDocument(clickAction3));
498 
499         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
500         mAppSearchManagerServiceStub.putDocuments(
501                 new PutDocumentsAidlRequest(
502                         AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
503                         DATABASE_NAME,
504                         new DocumentsParcel(
505                                 Collections.emptyList(), takenActionGenericDocumentParcels),
506                         mUserHandle,
507                         BINDER_CALL_START_TIME),
508                 callback);
509         assertThat(callback.get()).isNull(); // null means there wasn't an error
510         verifyCallStats(
511                 mContext.getPackageName(), DATABASE_NAME, CallStats.CALL_TYPE_PUT_DOCUMENTS);
512 
513         // Verify search sessions.
514         ArgumentCaptor<List<SearchSessionStats>> searchSessionsStatsCaptor =
515                 ArgumentCaptor.forClass(List.class);
516         verify(mLogger, timeout(1000).times(1)).logStats(searchSessionsStatsCaptor.capture());
517         List<SearchSessionStats> searchSessionsStats = searchSessionsStatsCaptor.getValue();
518 
519         assertThat(searchSessionsStats).hasSize(1);
520         assertThat(searchSessionsStats.get(0).getPackageName())
521                 .isEqualTo(mContext.getPackageName());
522         assertThat(searchSessionsStats.get(0).getDatabase()).isEqualTo(DATABASE_NAME);
523 
524         // Verify search intents.
525         List<SearchIntentStats> searchIntentsStats =
526                 searchSessionsStats.get(0).getSearchIntentsStats();
527         assertThat(searchIntentsStats).hasSize(2);
528 
529         assertThat(searchIntentsStats.get(0).getPackageName()).isEqualTo(mContext.getPackageName());
530         assertThat(searchIntentsStats.get(0).getDatabase()).isEqualTo(DATABASE_NAME);
531         assertThat(searchIntentsStats.get(0).getPrevQuery()).isNull();
532         assertThat(searchIntentsStats.get(0).getCurrQuery()).isEqualTo("tes");
533         assertThat(searchIntentsStats.get(0).getTimestampMillis()).isEqualTo(1000);
534         assertThat(searchIntentsStats.get(0).getNumResultsFetched()).isEqualTo(20);
535         assertThat(searchIntentsStats.get(0).getQueryCorrectionType())
536                 .isEqualTo(SearchIntentStats.QUERY_CORRECTION_TYPE_FIRST_QUERY);
537         assertThat(searchIntentsStats.get(0).getClicksStats()).hasSize(2);
538 
539         assertThat(searchIntentsStats.get(1).getPackageName()).isEqualTo(mContext.getPackageName());
540         assertThat(searchIntentsStats.get(1).getDatabase()).isEqualTo(DATABASE_NAME);
541         assertThat(searchIntentsStats.get(1).getPrevQuery()).isEqualTo("tes");
542         assertThat(searchIntentsStats.get(1).getCurrQuery()).isEqualTo("test");
543         assertThat(searchIntentsStats.get(1).getTimestampMillis()).isEqualTo(5000);
544         assertThat(searchIntentsStats.get(1).getNumResultsFetched()).isEqualTo(10);
545         assertThat(searchIntentsStats.get(1).getQueryCorrectionType())
546                 .isEqualTo(SearchIntentStats.QUERY_CORRECTION_TYPE_REFINEMENT);
547         assertThat(searchIntentsStats.get(1).getClicksStats()).hasSize(1);
548     }
549 
550     @Test
testLocalGetDocumentsStatsLogging()551     public void testLocalGetDocumentsStatsLogging() throws Exception {
552         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
553         mAppSearchManagerServiceStub.getDocuments(
554                 new GetDocumentsAidlRequest(
555                         AppSearchAttributionSource.createAttributionSource(mContext,
556                                 mCallingPid),
557                         mContext.getPackageName(), DATABASE_NAME,
558                         new GetByDocumentIdRequest.Builder(NAMESPACE)
559                                 .addIds(/* ids= */ Collections.emptyList())
560                                 .build(),
561                 mUserHandle, BINDER_CALL_START_TIME, /* isForEnterprise= */ false),
562                 callback);
563         assertThat(callback.get()).isNull(); // null means there wasn't an error
564         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
565                 CallStats.CALL_TYPE_GET_DOCUMENTS);
566     }
567 
568     @Test
testGlobalGetDocumentsStatsLogging()569     public void testGlobalGetDocumentsStatsLogging() throws Exception {
570         String otherPackageName = mContext.getPackageName() + "foo";
571         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
572         mAppSearchManagerServiceStub.getDocuments(
573                 new GetDocumentsAidlRequest(
574                         AppSearchAttributionSource.createAttributionSource(mContext,
575                                 mCallingPid),
576                         otherPackageName, DATABASE_NAME,
577                         new GetByDocumentIdRequest.Builder(NAMESPACE)
578                                 .addIds(/* ids= */ Collections.emptyList())
579                                 .build(),
580                        mUserHandle, BINDER_CALL_START_TIME, /* isForEnterprise= */ false),
581                 callback);
582         assertThat(callback.get()).isNull(); // null means there wasn't an error
583         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
584                 CallStats.CALL_TYPE_GLOBAL_GET_DOCUMENT_BY_ID);
585     }
586 
587     @Test
testSearchStatsLogging()588     public void testSearchStatsLogging() throws Exception {
589         TestResultCallback callback = new TestResultCallback();
590         mAppSearchManagerServiceStub.search(
591                 new SearchAidlRequest(AppSearchAttributionSource.createAttributionSource(mContext,
592                         mCallingPid),
593                         DATABASE_NAME, /* searchExpression= */ "", EMPTY_SEARCH_SPEC, mUserHandle,
594                         BINDER_CALL_START_TIME), callback);
595         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
596         verifyCallStats(mContext.getPackageName(), DATABASE_NAME, CallStats.CALL_TYPE_SEARCH);
597         // search only logs SearchStats indirectly so we don't verify it
598     }
599 
600     @Test
testGlobalSearchStatsLogging()601     public void testGlobalSearchStatsLogging() throws Exception {
602         TestResultCallback callback = new TestResultCallback();
603         mAppSearchManagerServiceStub.globalSearch(new GlobalSearchAidlRequest(
604                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
605                 /* searchExpression= */ "", EMPTY_SEARCH_SPEC, mUserHandle, BINDER_CALL_START_TIME,
606                 /* isForEnterprise= */ false), callback);
607         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
608         verifyCallStats(mContext.getPackageName(), CallStats.CALL_TYPE_GLOBAL_SEARCH);
609         // globalSearch only logs SearchStats indirectly so we don't verify it
610     }
611 
612     @Test
testLocalGetNextPageStatsLogging()613     public void testLocalGetNextPageStatsLogging() throws Exception {
614         TestResultCallback callback = new TestResultCallback();
615         mAppSearchManagerServiceStub.getNextPage(new GetNextPageAidlRequest(
616                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
617                 DATABASE_NAME, /* nextPageToken= */ 0,
618                 AppSearchSchema.StringPropertyConfig.JOINABLE_VALUE_TYPE_QUALIFIED_ID, mUserHandle,
619                 BINDER_CALL_START_TIME, /* isForEnterprise= */ false), callback);
620         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
621         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
622                 CallStats.CALL_TYPE_GET_NEXT_PAGE);
623         // getNextPage also logs SearchStats
624         ArgumentCaptor<SearchStats> searchStatsCaptor = ArgumentCaptor.forClass(SearchStats.class);
625         verify(mLogger, timeout(1000).times(1)).logStats(searchStatsCaptor.capture());
626         SearchStats searchStats = searchStatsCaptor.getValue();
627         assertThat(searchStats.getVisibilityScope()).isEqualTo(SearchStats.VISIBILITY_SCOPE_LOCAL);
628         assertThat(searchStats.getPackageName()).isEqualTo(mContext.getPackageName());
629         assertThat(searchStats.getDatabase()).isEqualTo(DATABASE_NAME);
630         assertThat(searchStats.getJoinType()).isEqualTo(
631                 AppSearchSchema.StringPropertyConfig.JOINABLE_VALUE_TYPE_QUALIFIED_ID);
632     }
633 
634     @Test
testGlobalGetNextPageStatsLogging()635     public void testGlobalGetNextPageStatsLogging() throws Exception {
636         TestResultCallback callback = new TestResultCallback();
637         mAppSearchManagerServiceStub.getNextPage(new GetNextPageAidlRequest(
638                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
639                 /* databaseName= */ null, /* nextPageToken= */ 0,
640                 AppSearchSchema.StringPropertyConfig.JOINABLE_VALUE_TYPE_QUALIFIED_ID, mUserHandle,
641                 BINDER_CALL_START_TIME, /* isForEnterprise= */ false), callback);
642         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
643         verifyCallStats(mContext.getPackageName(), CallStats.CALL_TYPE_GLOBAL_GET_NEXT_PAGE);
644         // getNextPage also logs SearchStats
645         ArgumentCaptor<SearchStats> searchStatsCaptor = ArgumentCaptor.forClass(SearchStats.class);
646         verify(mLogger, timeout(1000).times(1)).logStats(searchStatsCaptor.capture());
647         SearchStats searchStats = searchStatsCaptor.getValue();
648         assertThat(searchStats.getVisibilityScope()).isEqualTo(SearchStats.VISIBILITY_SCOPE_GLOBAL);
649         assertThat(searchStats.getPackageName()).isEqualTo(mContext.getPackageName());
650         assertThat(searchStats.getDatabase()).isNull();
651         assertThat(searchStats.getJoinType()).isEqualTo(
652                 AppSearchSchema.StringPropertyConfig.JOINABLE_VALUE_TYPE_QUALIFIED_ID);
653     }
654 
655     @Test
testInvalidateNextPageTokenStatsLogging()656     public void testInvalidateNextPageTokenStatsLogging() throws Exception {
657         mAppSearchManagerServiceStub.invalidateNextPageToken(new InvalidateNextPageTokenAidlRequest(
658                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
659                 /* nextPageToken= */ 0, mUserHandle, BINDER_CALL_START_TIME,
660                 /* isForEnterprise= */ false));
661         verifyCallStats(mContext.getPackageName(), CallStats.CALL_TYPE_INVALIDATE_NEXT_PAGE_TOKEN);
662     }
663 
664     @Test
testWriteSearchResultsToFileStatsLogging()665     public void testWriteSearchResultsToFileStatsLogging() throws Exception {
666         File tempFile = mTemporaryFolder.newFile();
667         FileDescriptor fd = IoBridge.open(tempFile.getPath(), O_WRONLY);
668         TestResultCallback callback = new TestResultCallback();
669         mAppSearchManagerServiceStub.writeSearchResultsToFile(
670                 new WriteSearchResultsToFileAidlRequest(
671                         AppSearchAttributionSource.createAttributionSource(mContext,
672                                 mCallingPid), DATABASE_NAME,
673                         new ParcelFileDescriptor(fd), /* searchExpression= */ "", EMPTY_SEARCH_SPEC,
674                         mUserHandle, BINDER_CALL_START_TIME), callback);
675         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
676         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
677                 CallStats.CALL_TYPE_WRITE_SEARCH_RESULTS_TO_FILE);
678     }
679 
680     @Test
testPutDocumentsFromFileStatsLogging()681     public void testPutDocumentsFromFileStatsLogging() throws Exception {
682         File tempFile = mTemporaryFolder.newFile();
683         FileDescriptor fd = IoBridge.open(tempFile.getPath(), O_RDONLY);
684         TestResultCallback callback = new TestResultCallback();
685         mAppSearchManagerServiceStub.putDocumentsFromFile(new PutDocumentsFromFileAidlRequest(
686                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
687                 DATABASE_NAME,
688                 new ParcelFileDescriptor(fd), mUserHandle,
689                 new SchemaMigrationStats.Builder(mContext.getPackageName(), DATABASE_NAME).build(),
690                 /* totalLatencyStartTimeMillis= */ 0, BINDER_CALL_START_TIME), callback);
691         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
692         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
693                 CallStats.CALL_TYPE_PUT_DOCUMENTS_FROM_FILE);
694         // putDocumentsFromFile also logs SchemaMigrationStats
695         ArgumentCaptor<SchemaMigrationStats> migrationStatsCaptor = ArgumentCaptor.forClass(
696                 SchemaMigrationStats.class);
697         verify(mLogger, timeout(1000).times(1)).logStats(migrationStatsCaptor.capture());
698         SchemaMigrationStats migrationStats = migrationStatsCaptor.getValue();
699         assertThat(migrationStats.getStatusCode()).isEqualTo(AppSearchResult.RESULT_OK);
700         assertThat(migrationStats.getSaveDocumentLatencyMillis()).isGreaterThan(0);
701         // putDocumentsFromFile only logs PutDocumentStats indirectly so we don't verify it
702     }
703 
704     @Test
testSearchSuggestionStatsLogging()705     public void testSearchSuggestionStatsLogging() throws Exception {
706         SearchSuggestionSpec searchSuggestionSpec =
707             new SearchSuggestionSpec.Builder(/*maximumResultCount=*/1).build();
708         TestResultCallback callback = new TestResultCallback();
709         mAppSearchManagerServiceStub.searchSuggestion(
710                 new SearchSuggestionAidlRequest(
711                         AppSearchAttributionSource.createAttributionSource(mContext,
712                                 mCallingPid),
713                         DATABASE_NAME, /* suggestionQueryExpression= */ "foo", searchSuggestionSpec,
714                         mUserHandle, BINDER_CALL_START_TIME),
715                 callback);
716         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
717         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
718                 CallStats.CALL_TYPE_SEARCH_SUGGESTION);
719     }
720 
721     @Test
testLocalReportUsageStatsLogging()722     public void testLocalReportUsageStatsLogging() throws Exception {
723         setUpTestSchema(mContext.getPackageName(), DATABASE_NAME);
724         setUpTestDocument(mContext.getPackageName(), DATABASE_NAME, NAMESPACE, ID);
725         TestResultCallback callback = new TestResultCallback();
726         mAppSearchManagerServiceStub.reportUsage(
727                 new ReportUsageAidlRequest(
728                         AppSearchAttributionSource.createAttributionSource(mContext,
729                                 mCallingPid),
730                         mContext.getPackageName(), DATABASE_NAME,
731                         new ReportUsageRequest.Builder(NAMESPACE, ID)
732                                 .setUsageTimestampMillis(/* usageTimestampMillis= */ 0)
733                                 .build(),
734                         /* systemUsage= */ false, mUserHandle, BINDER_CALL_START_TIME),
735                 callback);
736         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
737         verifyCallStats(mContext.getPackageName(), DATABASE_NAME, CallStats.CALL_TYPE_REPORT_USAGE);
738         removeTestSchema(mContext.getPackageName(), DATABASE_NAME);
739     }
740 
741     @Test
testGlobalReportUsageStatsLogging()742     public void testGlobalReportUsageStatsLogging() throws Exception {
743         // Grant system access for global report usage
744         mUiAutomation.adoptShellPermissionIdentity(Manifest.permission.READ_GLOBAL_APP_SEARCH_DATA);
745         try {
746             String otherPackageName = mContext.getPackageName() + "foo";
747             setUpTestSchema(otherPackageName, DATABASE_NAME);
748             setUpTestDocument(otherPackageName, DATABASE_NAME, NAMESPACE, ID);
749             TestResultCallback callback = new TestResultCallback();
750             mAppSearchManagerServiceStub.reportUsage(
751                     new ReportUsageAidlRequest(
752                             AppSearchAttributionSource.createAttributionSource(mContext,
753                                     mCallingPid),
754                             otherPackageName, DATABASE_NAME,
755                             new ReportUsageRequest.Builder(NAMESPACE, ID)
756                                     .setUsageTimestampMillis(/* usageTimestampMillis= */ 0)
757                                     .build(),
758                             /* systemUsage= */ true, mUserHandle, BINDER_CALL_START_TIME),
759                     callback);
760             assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
761             verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
762                     CallStats.CALL_TYPE_REPORT_SYSTEM_USAGE);
763             removeTestSchema(otherPackageName, DATABASE_NAME);
764         } finally {
765             mUiAutomation.dropShellPermissionIdentity();
766         }
767     }
768 
769     @Test
testRemoveByDocumentIdStatsLogging()770     public void testRemoveByDocumentIdStatsLogging() throws Exception {
771         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
772         mAppSearchManagerServiceStub.removeByDocumentId(
773                 new RemoveByDocumentIdAidlRequest(
774                         AppSearchAttributionSource.createAttributionSource(mContext,
775                                 mCallingPid),
776                         DATABASE_NAME,
777                         new RemoveByDocumentIdRequest.Builder(NAMESPACE)
778                                 .addIds(/* ids= */ Collections.emptyList())
779                                 .build(),
780                         mUserHandle,
781                         BINDER_CALL_START_TIME),
782                 callback);
783         assertThat(callback.get()).isNull(); // null means there wasn't an error
784         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
785                 CallStats.CALL_TYPE_REMOVE_DOCUMENTS_BY_ID);
786     }
787 
788     @Test
testRemoveByQueryStatsLogging()789     public void testRemoveByQueryStatsLogging() throws Exception {
790         TestResultCallback callback = new TestResultCallback();
791         mAppSearchManagerServiceStub.removeByQuery(
792                 new RemoveByQueryAidlRequest(
793                         AppSearchAttributionSource.createAttributionSource(mContext,
794                                 mCallingPid), DATABASE_NAME,
795                         /* queryExpression= */ "", EMPTY_SEARCH_SPEC, mUserHandle,
796                         BINDER_CALL_START_TIME),
797                 callback);
798         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
799         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
800                 CallStats.CALL_TYPE_REMOVE_DOCUMENTS_BY_SEARCH);
801     }
802 
803     @Test
testGetStorageInfoStatsLogging()804     public void testGetStorageInfoStatsLogging() throws Exception {
805         TestResultCallback callback = new TestResultCallback();
806         mAppSearchManagerServiceStub.getStorageInfo(
807                 new GetStorageInfoAidlRequest(
808                         AppSearchAttributionSource.createAttributionSource(mContext,
809                                 mCallingPid), DATABASE_NAME,
810                         mUserHandle, BINDER_CALL_START_TIME),
811                 callback);
812         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
813         verifyCallStats(mContext.getPackageName(), DATABASE_NAME,
814                 CallStats.CALL_TYPE_GET_STORAGE_INFO);
815     }
816 
817     @Test
testPersistToDiskStatsLogging()818     public void testPersistToDiskStatsLogging() throws Exception {
819         mAppSearchManagerServiceStub.persistToDisk(
820                 new PersistToDiskAidlRequest(
821                         AppSearchAttributionSource.createAttributionSource(mContext,
822                                 mCallingPid), mUserHandle,
823                         BINDER_CALL_START_TIME));
824         verifyCallStats(mContext.getPackageName(), CallStats.CALL_TYPE_FLUSH);
825     }
826 
827     @Test
testRegisterObserverCallbackStatsLogging()828     public void testRegisterObserverCallbackStatsLogging() throws Exception {
829         AppSearchResultParcel<Void> resultParcel =
830                 mAppSearchManagerServiceStub.registerObserverCallback(
831                         new RegisterObserverCallbackAidlRequest(
832                                 AppSearchAttributionSource.createAttributionSource(mContext,
833                                         mCallingPid),
834                                 mContext.getPackageName(),
835                                 new ObserverSpec.Builder().build(),
836                                 mUserHandle, BINDER_CALL_START_TIME),
837                         new IAppSearchObserverProxy.Stub() {
838                             @Override
839                             public void onSchemaChanged(String packageName, String databaseName,
840                                     List<String> changedSchemaNames) throws RemoteException {
841                             }
842 
843                             @Override
844                             public void onDocumentChanged(String packageName, String databaseName,
845                                     String namespace, String schemaName,
846                                     List<String> changedDocumentIds) throws RemoteException {
847                             }
848                         });
849         assertThat(resultParcel.getResult().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
850         verifyCallStats(mContext.getPackageName(), CallStats.CALL_TYPE_REGISTER_OBSERVER_CALLBACK);
851     }
852 
853     @Test
testUnregisterObserverCallbackStatsLogging()854     public void testUnregisterObserverCallbackStatsLogging() throws Exception {
855         AppSearchResultParcel<Void> resultParcel =
856                 mAppSearchManagerServiceStub.unregisterObserverCallback(
857                         new UnregisterObserverCallbackAidlRequest(
858                                 AppSearchAttributionSource.createAttributionSource(mContext,
859                                         mCallingPid),
860                                 mContext.getPackageName(), mUserHandle,
861                                 BINDER_CALL_START_TIME),
862                         new IAppSearchObserverProxy.Stub() {
863                             @Override
864                             public void onSchemaChanged(String packageName, String databaseName,
865                                     List<String> changedSchemaNames) throws RemoteException {
866                             }
867 
868                             @Override
869                             public void onDocumentChanged(String packageName, String databaseName,
870                                     String namespace, String schemaName,
871                                     List<String> changedDocumentIds) throws RemoteException {
872                             }
873                         });
874         assertThat(resultParcel.getResult().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
875         verifyCallStats(mContext.getPackageName(),
876                 CallStats.CALL_TYPE_UNREGISTER_OBSERVER_CALLBACK);
877     }
878 
879     @Test
testInitializeStatsLogging()880     public void testInitializeStatsLogging() throws Exception {
881         TestResultCallback callback = new TestResultCallback();
882         mAppSearchManagerServiceStub.initialize(
883                 new InitializeAidlRequest(
884                         AppSearchAttributionSource.createAttributionSource(mContext,
885                                 mCallingPid), mUserHandle,
886                         BINDER_CALL_START_TIME),
887                 callback);
888         assertThat(callback.get().getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
889         verifyCallStats(mContext.getPackageName(), CallStats.CALL_TYPE_INITIALIZE);
890         // initialize only logs InitializeStats indirectly so we don't verify it
891     }
892 
verifyCallStats(String packageName, String databaseName, int callType)893     private void verifyCallStats(String packageName, String databaseName, int callType) {
894         ArgumentCaptor<CallStats> captor = ArgumentCaptor.forClass(CallStats.class);
895         verify(mLogger, timeout(1000).times(1)).logStats(captor.capture());
896         CallStats callStats = captor.getValue();
897         assertThat(callStats.getPackageName()).isEqualTo(packageName);
898         assertThat(callStats.getDatabase()).isEqualTo(databaseName);
899         assertThat(callStats.getCallType()).isEqualTo(callType);
900         assertThat(callStats.getStatusCode()).isEqualTo(AppSearchResult.RESULT_OK);
901         assertThat(callStats.getEstimatedBinderLatencyMillis()).isGreaterThan(0);
902     }
903 
verifyCallStats(String packageName, int callType)904     private void verifyCallStats(String packageName, int callType) {
905         verifyCallStats(packageName, /* databaseName= */ null, callType);
906     }
907 
908     @Test
testDenylistMatchingCallingPackage()909     public void testDenylistMatchingCallingPackage() throws Exception {
910         String denylistString =
911                 "pkg=com.android.appsearch.mockingservicestests&apis=localSetSchema,"
912                         + "globalGetSchema,localGetSchema,localGetNamespaces,localPutDocuments,"
913                         + "globalGetDocuments,localGetDocuments,localSearch,globalSearch,"
914                         + "globalGetNextPage,localGetNextPage,invalidateNextPageToken,"
915                         + "localWriteSearchResultsToFile,localPutDocumentsFromFile,"
916                         + "localSearchSuggestion,globalReportUsage,localReportUsage,"
917                         + "localRemoveByDocumentId,localRemoveBySearch,localGetStorageInfo,flush,"
918                         + "globalRegisterObserverCallback,globalUnregisterObserverCallback,"
919                         + "initialize,executeAppFunction";
920         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
921                 KEY_DENYLIST, denylistString, false);
922         // We expect all local calls (pkg+db) and global calls (pkg only) to be denied since the
923         // denylist denies all api's for our calling package
924         verifyLocalCallsResults(RESULT_DENIED);
925         verifyGlobalCallsResults(RESULT_DENIED);
926     }
927 
928     @Test
testDenylistNonMatchingCallingPackage()929     public void testDenylistNonMatchingCallingPackage() throws Exception {
930         String denylistString =
931                 "pkg=foo&apis=localSetSchema,globalGetSchema,localGetSchema,localGetNamespaces,"
932                         + "localPutDocuments,globalGetDocuments,localGetDocuments,localSearch,"
933                         + "globalSearch,globalGetNextPage,localGetNextPage,invalidateNextPageToken,"
934                         + "localWriteSearchResultsToFile,localPutDocumentsFromFile,"
935                         + "localSearchSuggestion,globalReportUsage,localReportUsage,"
936                         + "localRemoveByDocumentId,localRemoveBySearch,localGetStorageInfo,flush,"
937                         + "globalRegisterObserverCallback,globalUnregisterObserverCallback,"
938                         + "initialize,executeAppFunction";
939         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
940                 KEY_DENYLIST, denylistString, false);
941         // We expect none of the local calls (pkg+db) and global calls (pkg only) to be denied since
942         // the denylist denies all api's for a different calling package
943         verifyLocalCallsResults(AppSearchResult.RESULT_OK);
944         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
945     }
946 
947     @Test
testDenylistMatchingCallingPackageAndDatabase()948     public void testDenylistMatchingCallingPackageAndDatabase() throws Exception {
949         String denylistString =
950                 "pkg=com.android.appsearch.mockingservicestests&db=databaseName&apis="
951                         + "localSetSchema,globalGetSchema,localGetSchema,localGetNamespaces,"
952                         + "localPutDocuments,globalGetDocuments,localGetDocuments,localSearch,"
953                         + "globalSearch,globalGetNextPage,localGetNextPage,invalidateNextPageToken,"
954                         + "localWriteSearchResultsToFile,localPutDocumentsFromFile,"
955                         + "localSearchSuggestion,globalReportUsage,localReportUsage,"
956                         + "localRemoveByDocumentId,localRemoveBySearch,localGetStorageInfo,flush,"
957                         + "globalRegisterObserverCallback,globalUnregisterObserverCallback,"
958                         + "initialize,executeAppFunction";
959         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
960                 KEY_DENYLIST, denylistString, false);
961         // We expect only the local calls (pkg+db) to be denied since the denylist specifies a
962         // package-database name pair
963         verifyLocalCallsResults(RESULT_DENIED);
964         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
965     }
966 
967     @Test
testDenylistNonMatchingCallingPackageAndDatabase()968     public void testDenylistNonMatchingCallingPackageAndDatabase() throws Exception {
969         // This denylist has two entries both of which should not match any of the api calls below
970         // since either the package name or the database name will be different
971         String denylistString =
972                 "pkg=foo&db=databaseName&apis=localSetSchema,globalGetSchema,localGetSchema,"
973                         + "localGetNamespaces,localPutDocuments,globalGetDocuments,"
974                         + "localGetDocuments,localSearch,globalSearch,globalGetNextPage,"
975                         + "localGetNextPage,invalidateNextPageToken,localWriteSearchResultsToFile,"
976                         + "localPutDocumentsFromFile,localSearchSuggestion,globalReportUsage,"
977                         + "localReportUsage,localRemoveByDocumentId,localRemoveBySearch,"
978                         + "localGetStorageInfo,flush,globalRegisterObserverCallback,"
979                         + "globalUnregisterObserverCallback,initialize;"
980                         + "pkg=com.android.appsearch.mockingservicestests&db=foo&apis="
981                         + "localSetSchema,globalGetSchema,localGetSchema,localGetNamespaces,"
982                         + "localPutDocuments,globalGetDocuments,localGetDocuments,localSearch,"
983                         + "globalSearch,globalGetNextPage,localGetNextPage,invalidateNextPageToken,"
984                         + "localWriteSearchResultsToFile,localPutDocumentsFromFile,"
985                         + "localSearchSuggestion,globalReportUsage,localReportUsage,"
986                         + "localRemoveByDocumentId,localRemoveBySearch,localGetStorageInfo,flush,"
987                         + "globalRegisterObserverCallback,globalUnregisterObserverCallback,"
988                         + "initialize,executeAppFunction";
989         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
990                 KEY_DENYLIST, denylistString, false);
991         verifyLocalCallsResults(AppSearchResult.RESULT_OK);
992         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
993     }
994 
995     @Test
testDenylistMatchingCallingDatabase()996     public void testDenylistMatchingCallingDatabase() throws Exception {
997         String denylistString =
998                 "db=databaseName&apis=localSetSchema,globalGetSchema,localGetSchema,"
999                         + "localGetNamespaces,localPutDocuments,globalGetDocuments,"
1000                         + "localGetDocuments,localSearch,globalSearch,globalGetNextPage,"
1001                         + "localGetNextPage,invalidateNextPageToken,localWriteSearchResultsToFile,"
1002                         + "localPutDocumentsFromFile,localSearchSuggestion,globalReportUsage,"
1003                         + "localReportUsage,localRemoveByDocumentId,localRemoveBySearch,"
1004                         + "localGetStorageInfo,flush,globalRegisterObserverCallback,"
1005                         + "globalUnregisterObserverCallback,initialize,executeAppFunction";
1006         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1007                 KEY_DENYLIST, denylistString, false);
1008 
1009         verifyLocalCallsResults(RESULT_DENIED);
1010         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1011 
1012         // Add mocking to spy'd package manager to return current uid for package foo
1013         // This is necessary to pass call verification using a different package name
1014         PackageManager spyPackageManager = mContext.getPackageManager();
1015         int uid = AppSearchAttributionSource.createAttributionSource(mContext,
1016                 mCallingPid).getUid();
1017         doReturn(uid).when(spyPackageManager).getPackageUid(FOO_PACKAGE_NAME, /* flags= */ 0);
1018         // Specifically grant permission for report system usage to package foo
1019         doReturn(PackageManager.PERMISSION_GRANTED).when(spyPackageManager).checkPermission(
1020                 READ_GLOBAL_APP_SEARCH_DATA, FOO_PACKAGE_NAME);
1021 
1022         // Change the calling package name used in the helper methods indirectly through a newly
1023         // wrapped context
1024         Context context = ApplicationProvider.getApplicationContext();
1025         mContext = new ContextWrapper(context) {
1026             @Override
1027             public String getPackageName() {
1028                 return FOO_PACKAGE_NAME;
1029             }
1030 
1031             @Override
1032             public AttributionSource getAttributionSource() {
1033                 return super.getAttributionSource().withPackageName(FOO_PACKAGE_NAME);
1034             }
1035         };
1036 
1037         // Confirm that we're using a different package name
1038         assertThat(mContext.getPackageName()).isEqualTo(FOO_PACKAGE_NAME);
1039         assertThat(AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid)
1040                 .getPackageName()).isEqualTo(FOO_PACKAGE_NAME);
1041 
1042         verifyLocalCallsResults(RESULT_DENIED);
1043         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1044     }
1045 
1046     @Test
testDenylistNonMatchingCallingDatabase()1047     public void testDenylistNonMatchingCallingDatabase() throws Exception {
1048         String denylistString =
1049                 "db=foo&apis=localSetSchema,globalGetSchema,localGetSchema,localGetNamespaces,"
1050                         + "localPutDocuments,globalGetDocuments,localGetDocuments,localSearch,"
1051                         + "globalSearch,globalGetNextPage,localGetNextPage,invalidateNextPageToken,"
1052                         + "localWriteSearchResultsToFile,localPutDocumentsFromFile,"
1053                         + "localSearchSuggestion,globalReportUsage,localReportUsage,"
1054                         + "localRemoveByDocumentId,localRemoveBySearch,localGetStorageInfo,flush,"
1055                         + "globalRegisterObserverCallback,globalUnregisterObserverCallback,"
1056                         + "initialize,executeAppFunction";
1057         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1058                 KEY_DENYLIST, denylistString, false);
1059         verifyLocalCallsResults(AppSearchResult.RESULT_OK);
1060         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1061     }
1062 
1063     @Test
testDenylistSomeApis()1064     public void testDenylistSomeApis() throws Exception {
1065         String denylistString =
1066                 "pkg=com.android.appsearch.mockingservicestests&apis=localSetSchema,localGetSchema,"
1067                         + "localGetNamespaces";
1068         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1069                 KEY_DENYLIST, denylistString, false);
1070         // Specified APIs
1071         verifySetSchemaResult(RESULT_DENIED);
1072         verifyLocalGetSchemaResult(RESULT_DENIED);
1073         verifyGetNamespacesResult(RESULT_DENIED);
1074         // Everything else
1075         verifyPutDocumentsResult(AppSearchResult.RESULT_OK);
1076         verifyLocalGetDocumentsResult(AppSearchResult.RESULT_OK);
1077         verifySearchResult(AppSearchResult.RESULT_OK);
1078         verifyGlobalGetSchemaResult(AppSearchResult.RESULT_OK);
1079         verifyLocalGetNextPageResult(AppSearchResult.RESULT_OK);
1080         verifyWriteSearchResultsToFileResult(AppSearchResult.RESULT_OK);
1081         verifyPutDocumentsFromFileResult(AppSearchResult.RESULT_OK);
1082         verifySearchSuggestionResult(AppSearchResult.RESULT_OK);
1083         verifyLocalReportUsageResult(AppSearchResult.RESULT_OK);
1084         verifyRemoveByDocumentIdResult(AppSearchResult.RESULT_OK);
1085         verifyRemoveByQueryResult(AppSearchResult.RESULT_OK);
1086         verifyGetStorageInfoResult(AppSearchResult.RESULT_OK);
1087         verifyGlobalGetDocumentsResult(AppSearchResult.RESULT_OK);
1088         verifyGlobalSearchResult(AppSearchResult.RESULT_OK);
1089         verifyGlobalGetNextPageResult(AppSearchResult.RESULT_OK);
1090         verifyInvalidateNextPageTokenResult(AppSearchResult.RESULT_OK);
1091         verifyGlobalReportUsageResult(AppSearchResult.RESULT_OK);
1092         verifyPersistToDiskResult(AppSearchResult.RESULT_OK);
1093         verifyRegisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1094         verifyUnregisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1095         verifyInitializeResult(AppSearchResult.RESULT_OK);
1096         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_OK);
1097     }
1098 
1099     @Test
testAppSearchRateLimit_rateLimitOn_allApis()1100     public void testAppSearchRateLimit_rateLimitOn_allApis() throws Exception {
1101         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1102                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(true), false);
1103         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1104                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(1), false);
1105         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1106                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1107                 Float.toString(0.8f),
1108                 false);
1109         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1110                 KEY_RATE_LIMIT_API_COSTS,
1111                 "localSetSchema:5;globalGetSchema:5;localGetSchema:5;localGetNamespaces:5;"
1112                         + "localPutDocuments:5;globalGetDocuments:5;localGetDocuments:5;"
1113                         + "localSearch:5;globalSearch:5;globalGetNextPage:5;localGetNextPage:5;"
1114                         + "invalidateNextPageToken:5;localWriteSearchResultsToFile:5;"
1115                         + "localPutDocumentsFromFile:5;localSearchSuggestion:5;"
1116                         + "globalReportUsage:5;localReportUsage:5;localRemoveByDocumentId:5;"
1117                         + "localRemoveBySearch:5;localGetStorageInfo:5;flush:5;"
1118                         + "executeAppFunction:5",
1119                 false);
1120 
1121         verifySetSchemaResult(RESULT_RATE_LIMITED);
1122         verifyLocalGetSchemaResult(RESULT_RATE_LIMITED);
1123         verifySearchResult(RESULT_RATE_LIMITED);
1124         verifyPutDocumentsResult(RESULT_RATE_LIMITED);
1125         verifyLocalGetDocumentsResult(RESULT_RATE_LIMITED);
1126         verifyLocalGetNextPageResult(RESULT_RATE_LIMITED);
1127         verifyGlobalGetDocumentsResult(RESULT_RATE_LIMITED);
1128         verifyGlobalSearchResult(RESULT_RATE_LIMITED);
1129         verifyGlobalGetNextPageResult(RESULT_RATE_LIMITED);
1130         verifyInvalidateNextPageTokenResult(RESULT_RATE_LIMITED);
1131         verifyGlobalReportUsageResult(RESULT_RATE_LIMITED);
1132         verifyPersistToDiskResult(RESULT_RATE_LIMITED);
1133         verifyExecuteAppFunctionCallbackResult(RESULT_RATE_LIMITED);
1134 
1135         // initialize, registerObserver and unregisterObserver do not have rate limit.
1136     }
1137 
1138     @Test
testAppSearchRateLimit_rateLimitOff_acceptTasksAndIgnoreCapacities()1139     public void testAppSearchRateLimit_rateLimitOff_acceptTasksAndIgnoreCapacities()
1140             throws Exception {
1141         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1142                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(false), false);
1143         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1144                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(10), false);
1145         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1146                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1147                 Float.toString(0.8f),
1148                 false);
1149         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1150                 KEY_RATE_LIMIT_API_COSTS,
1151                 "localSearch:6;localSetSchema:9;localGetSchema:15",
1152                 false);
1153 
1154         // All rate limits should be ignored when rate limit is off
1155         verifyLocalCallsResults(AppSearchResult.RESULT_OK);
1156         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1157     }
1158 
1159     @Test
testAppSearchRateLimit_rateLimitOn_dropTaskDueToCapacitiesExceeded()1160     public void testAppSearchRateLimit_rateLimitOn_dropTaskDueToCapacitiesExceeded()
1161             throws Exception {
1162         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1163                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(true), false);
1164         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1165                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(10), false);
1166         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1167                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1168                 Float.toString(0.8f),
1169                 false);
1170         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1171                 KEY_RATE_LIMIT_API_COSTS,
1172                 "localSearch:6;localSetSchema:9;localGetSchema:15",
1173                 false);
1174 
1175         // Set Schema call is rejected because of per-package capacity exceeded
1176         verifySetSchemaResult(RESULT_RATE_LIMITED);
1177         // Set Schema call is rejected because of total capacity exceeded
1178         verifyLocalGetSchemaResult(RESULT_RATE_LIMITED);
1179         // Other calls should be fine
1180         verifySearchResult(AppSearchResult.RESULT_OK);
1181         verifyPutDocumentsResult(AppSearchResult.RESULT_OK);
1182         verifyLocalGetDocumentsResult(AppSearchResult.RESULT_OK);
1183         verifyLocalGetNextPageResult(AppSearchResult.RESULT_OK);
1184         verifyGlobalGetDocumentsResult(AppSearchResult.RESULT_OK);
1185         verifyGlobalSearchResult(AppSearchResult.RESULT_OK);
1186         verifyGlobalGetNextPageResult(AppSearchResult.RESULT_OK);
1187         verifyInvalidateNextPageTokenResult(AppSearchResult.RESULT_OK);
1188         verifyGlobalReportUsageResult(AppSearchResult.RESULT_OK);
1189         verifyPersistToDiskResult(AppSearchResult.RESULT_OK);
1190         verifyRegisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1191         verifyUnregisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1192         verifyInitializeResult(AppSearchResult.RESULT_OK);
1193         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_OK);
1194     }
1195 
1196     @Test
testAppSearchRateLimit_rateLimitOn_noTasksDropped()1197     public void testAppSearchRateLimit_rateLimitOn_noTasksDropped() throws Exception {
1198         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1199                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(true), false);
1200         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1201                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(1000), false);
1202         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1203                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1204                 Float.toString(0.8f),
1205                 false);
1206         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1207                 KEY_RATE_LIMIT_API_COSTS,
1208                 "localSearch:6;localSetSchema:9;localGetSchema:15",
1209                 false);
1210 
1211         verifyLocalCallsResults(AppSearchResult.RESULT_OK);
1212         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1213     }
1214 
1215     @Test
testAppSearchRateLimit_rateLimitOnToOff()1216     public void testAppSearchRateLimit_rateLimitOnToOff() throws Exception {
1217         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1218                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(true), false);
1219         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1220                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(10), false);
1221         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1222                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1223                 Float.toString(0.8f),
1224                 false);
1225         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1226                 KEY_RATE_LIMIT_API_COSTS,
1227                 "localSearch:6;localSetSchema:9;localGetSchema:15",
1228                 false);
1229         verifySetSchemaResult(RESULT_RATE_LIMITED);
1230         verifyLocalGetSchemaResult(RESULT_RATE_LIMITED);
1231         verifySearchResult(AppSearchResult.RESULT_OK);
1232         verifyPutDocumentsResult(AppSearchResult.RESULT_OK);
1233         verifyLocalGetDocumentsResult(AppSearchResult.RESULT_OK);
1234         verifyLocalGetNextPageResult(AppSearchResult.RESULT_OK);
1235         verifyGlobalGetDocumentsResult(AppSearchResult.RESULT_OK);
1236         verifyGlobalSearchResult(AppSearchResult.RESULT_OK);
1237         verifyGlobalGetNextPageResult(AppSearchResult.RESULT_OK);
1238         verifyInvalidateNextPageTokenResult(AppSearchResult.RESULT_OK);
1239         verifyGlobalReportUsageResult(AppSearchResult.RESULT_OK);
1240         verifyPersistToDiskResult(AppSearchResult.RESULT_OK);
1241         verifyRegisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1242         verifyUnregisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1243         verifyInitializeResult(AppSearchResult.RESULT_OK);
1244         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_OK);
1245 
1246         // All calls should be fine after switching rate limiting to off
1247         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1248                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(false), false);
1249         verifyLocalCallsResults(AppSearchResult.RESULT_OK);
1250         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1251     }
1252 
1253     @Test
testAppSearchRateLimit_rateLimitOffToOn()1254     public void testAppSearchRateLimit_rateLimitOffToOn() throws Exception {
1255         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1256                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(false), false);
1257         verifyLocalCallsResults(AppSearchResult.RESULT_OK);
1258         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1259 
1260         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1261                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(true), false);
1262         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1263                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(5), false);
1264         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1265                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1266                 Float.toString(0.8f),
1267                 false);
1268         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1269                 KEY_RATE_LIMIT_API_COSTS,
1270                 "localSearch:6;localSetSchema:9;localGetSchema:15",
1271                 false);
1272         // Some calls are rejected once rate limiting gets enabled
1273         verifySearchResult(RESULT_RATE_LIMITED);
1274         verifySetSchemaResult(RESULT_RATE_LIMITED);
1275         verifyLocalGetSchemaResult(RESULT_RATE_LIMITED);
1276         verifyPutDocumentsResult(AppSearchResult.RESULT_OK);
1277         verifyLocalGetDocumentsResult(AppSearchResult.RESULT_OK);
1278         verifyLocalGetNextPageResult(AppSearchResult.RESULT_OK);
1279         verifyGlobalGetDocumentsResult(AppSearchResult.RESULT_OK);
1280         verifyGlobalSearchResult(AppSearchResult.RESULT_OK);
1281         verifyGlobalGetNextPageResult(AppSearchResult.RESULT_OK);
1282         verifyInvalidateNextPageTokenResult(AppSearchResult.RESULT_OK);
1283         verifyGlobalReportUsageResult(AppSearchResult.RESULT_OK);
1284         verifyPersistToDiskResult(AppSearchResult.RESULT_OK);
1285         verifyRegisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1286         verifyUnregisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1287         verifyInitializeResult(AppSearchResult.RESULT_OK);
1288         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_OK);
1289     }
1290 
1291     @Test
testAppSearchRateLimit_rateLimitChangeToHigherLimit()1292     public void testAppSearchRateLimit_rateLimitChangeToHigherLimit() throws Exception {
1293         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1294                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(true), false);
1295         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1296                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(10), false);
1297         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1298                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1299                 Float.toString(0.8f),
1300                 false);
1301         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1302                 KEY_RATE_LIMIT_API_COSTS,
1303                 "localSearch:6;localSetSchema:9;localGetSchema:15",
1304                 false);
1305         verifySetSchemaResult(RESULT_RATE_LIMITED);
1306         verifyLocalGetSchemaResult(RESULT_RATE_LIMITED);
1307 
1308         verifySearchResult(AppSearchResult.RESULT_OK);
1309         verifyPutDocumentsResult(AppSearchResult.RESULT_OK);
1310         verifyLocalGetDocumentsResult(AppSearchResult.RESULT_OK);
1311         verifyLocalGetNextPageResult(AppSearchResult.RESULT_OK);
1312         verifyGlobalGetDocumentsResult(AppSearchResult.RESULT_OK);
1313         verifyGlobalSearchResult(AppSearchResult.RESULT_OK);
1314         verifyGlobalGetNextPageResult(AppSearchResult.RESULT_OK);
1315         verifyInvalidateNextPageTokenResult(AppSearchResult.RESULT_OK);
1316         verifyGlobalReportUsageResult(AppSearchResult.RESULT_OK);
1317         verifyPersistToDiskResult(AppSearchResult.RESULT_OK);
1318         verifyRegisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1319         verifyUnregisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1320         verifyInitializeResult(AppSearchResult.RESULT_OK);
1321         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1322 
1323         // Only getSchema call should be rejected after setting to higher limit
1324         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1325                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(15), false);
1326         verifyLocalGetSchemaResult(RESULT_RATE_LIMITED);
1327 
1328         verifySetSchemaResult(AppSearchResult.RESULT_OK);
1329         verifySearchResult(AppSearchResult.RESULT_OK);
1330         verifyPutDocumentsResult(AppSearchResult.RESULT_OK);
1331         verifyLocalGetDocumentsResult(AppSearchResult.RESULT_OK);
1332         verifyLocalGetNextPageResult(AppSearchResult.RESULT_OK);
1333         verifyGlobalGetDocumentsResult(AppSearchResult.RESULT_OK);
1334         verifyGlobalSearchResult(AppSearchResult.RESULT_OK);
1335         verifyGlobalGetNextPageResult(AppSearchResult.RESULT_OK);
1336         verifyInvalidateNextPageTokenResult(AppSearchResult.RESULT_OK);
1337         verifyGlobalReportUsageResult(AppSearchResult.RESULT_OK);
1338         verifyPersistToDiskResult(AppSearchResult.RESULT_OK);
1339         verifyRegisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1340         verifyUnregisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1341         verifyInitializeResult(AppSearchResult.RESULT_OK);
1342         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1343     }
1344 
1345     @Test
testAppSearchRateLimit_rateLimitChangeToLowerLimit()1346     public void testAppSearchRateLimit_rateLimitChangeToLowerLimit() throws Exception {
1347         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1348                 KEY_RATE_LIMIT_ENABLED, Boolean.toString(true), false);
1349         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1350                 KEY_RATE_LIMIT_TASK_QUEUE_TOTAL_CAPACITY, Integer.toString(10), false);
1351         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1352                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1353                 Float.toString(0.8f),
1354                 false);
1355         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1356                 KEY_RATE_LIMIT_API_COSTS,
1357                 "localSearch:6;localSetSchema:9;localGetSchema:15",
1358                 false);
1359         verifySetSchemaResult(RESULT_RATE_LIMITED);
1360         verifyLocalGetSchemaResult(RESULT_RATE_LIMITED);
1361 
1362         verifySearchResult(AppSearchResult.RESULT_OK);
1363         verifyPutDocumentsResult(AppSearchResult.RESULT_OK);
1364         verifyLocalGetDocumentsResult(AppSearchResult.RESULT_OK);
1365         verifyLocalGetNextPageResult(AppSearchResult.RESULT_OK);
1366         verifyGlobalGetDocumentsResult(AppSearchResult.RESULT_OK);
1367         verifyGlobalSearchResult(AppSearchResult.RESULT_OK);
1368         verifyGlobalGetNextPageResult(AppSearchResult.RESULT_OK);
1369         verifyInvalidateNextPageTokenResult(AppSearchResult.RESULT_OK);
1370         verifyGlobalReportUsageResult(AppSearchResult.RESULT_OK);
1371         verifyPersistToDiskResult(AppSearchResult.RESULT_OK);
1372         verifyRegisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1373         verifyUnregisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1374         verifyInitializeResult(AppSearchResult.RESULT_OK);
1375         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1376 
1377         // Search call should also get rejected after setting a lower limit
1378         DeviceConfig.setProperty(DeviceConfig.NAMESPACE_APPSEARCH,
1379                 KEY_RATE_LIMIT_TASK_QUEUE_PER_PACKAGE_CAPACITY_PERCENTAGE,
1380                 Float.toString(0.5f),
1381                 false);
1382         verifySearchResult(RESULT_RATE_LIMITED);
1383         verifySetSchemaResult(RESULT_RATE_LIMITED);
1384         verifyLocalGetSchemaResult(RESULT_RATE_LIMITED);
1385 
1386         verifyPutDocumentsResult(AppSearchResult.RESULT_OK);
1387         verifyLocalGetDocumentsResult(AppSearchResult.RESULT_OK);
1388         verifyLocalGetNextPageResult(AppSearchResult.RESULT_OK);
1389         verifyGlobalGetDocumentsResult(AppSearchResult.RESULT_OK);
1390         verifyGlobalSearchResult(AppSearchResult.RESULT_OK);
1391         verifyGlobalGetNextPageResult(AppSearchResult.RESULT_OK);
1392         verifyInvalidateNextPageTokenResult(AppSearchResult.RESULT_OK);
1393         verifyGlobalReportUsageResult(AppSearchResult.RESULT_OK);
1394         verifyPersistToDiskResult(AppSearchResult.RESULT_OK);
1395         verifyRegisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1396         verifyUnregisterObserverCallbackResult(AppSearchResult.RESULT_OK);
1397         verifyInitializeResult(AppSearchResult.RESULT_OK);
1398         verifyGlobalCallsResults(AppSearchResult.RESULT_OK);
1399     }
1400 
1401     @Test
testEnterpriseGetSchema_noEnterpriseUser_emptyResult()1402     public void testEnterpriseGetSchema_noEnterpriseUser_emptyResult() throws Exception {
1403         // Even on devices with an enterprise user, this test will run properly, since we haven't
1404         // unlocked the enterprise user for our local instance of AppSearchManagerService
1405         TestResultCallback callback = new TestResultCallback();
1406         mAppSearchManagerServiceStub.getSchema(
1407                 new GetSchemaAidlRequest(
1408                         AppSearchAttributionSource.createAttributionSource(mContext,
1409                                 mCallingPid),
1410                         mContext.getPackageName(), DATABASE_NAME, mUserHandle,
1411                         BINDER_CALL_START_TIME, /* isForEnterprise= */ true),
1412                 callback);
1413         AppSearchResult<GetSchemaResponse> result =
1414                 (AppSearchResult<GetSchemaResponse>) callback.get();
1415         assertThat(result.getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
1416         assertThat(result.getResultValue().getSchemas()).isEmpty();
1417         // No CallStats logged since we returned early
1418         verify(mLogger, timeout(1000).times(0)).logStats(any(CallStats.class));
1419     }
1420 
1421     @Test
testEnterpriseGetDocuments_noEnterpriseUser_emptyResult()1422     public void testEnterpriseGetDocuments_noEnterpriseUser_emptyResult() throws Exception {
1423         // Even on devices with an enterprise user, this test will run properly, since we haven't
1424         // unlocked the enterprise user for our local instance of AppSearchManagerService
1425         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
1426         mAppSearchManagerServiceStub.getDocuments(
1427                 new GetDocumentsAidlRequest(
1428                         AppSearchAttributionSource.createAttributionSource(mContext,
1429                                 mCallingPid),
1430                 mContext.getPackageName(), DATABASE_NAME, new GetByDocumentIdRequest.Builder(
1431                         NAMESPACE)
1432                         .addIds(/* ids= */ Collections.emptyList())
1433                         .build(),
1434                 mUserHandle, BINDER_CALL_START_TIME, /* isForEnterprise= */ true),
1435                 callback);
1436         assertThat(callback.get()).isNull(); // null means there wasn't an error
1437         assertThat(callback.getBatchResult().getAll()).isEmpty();
1438         // No CallStats logged since we returned early
1439         verify(mLogger, timeout(1000).times(0)).logStats(any(CallStats.class));
1440     }
1441 
1442     @Test
testEnterpriseGlobalSearch_noEnterpriseUser_emptyResult()1443     public void testEnterpriseGlobalSearch_noEnterpriseUser_emptyResult() throws Exception {
1444         // Even on devices with an enterprise user, this test will run properly, since we haven't
1445         // unlocked the enterprise user for our local instance of AppSearchManagerService
1446         TestResultCallback callback = new TestResultCallback();
1447         mAppSearchManagerServiceStub.globalSearch(new GlobalSearchAidlRequest(
1448                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
1449                 /* searchExpression= */ "", EMPTY_SEARCH_SPEC, mUserHandle, BINDER_CALL_START_TIME,
1450                 /* isForEnterprise= */ true), callback);
1451         AppSearchResult<SearchResultPage> result =
1452                 (AppSearchResult<SearchResultPage>) callback.get();
1453         assertThat(result.getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
1454         assertThat(result.getResultValue().getResults()).isEmpty();
1455         // No CallStats logged since we returned early
1456         verify(mLogger, timeout(1000).times(0)).logStats(any(CallStats.class));
1457     }
1458 
1459     @Test
testEnterpriseGetNextPage_noEnterpriseUser_emptyResult()1460     public void testEnterpriseGetNextPage_noEnterpriseUser_emptyResult() throws Exception {
1461         // Even on devices with an enterprise user, this test will run properly, since we haven't
1462         // unlocked the enterprise user for our local instance of AppSearchManagerService
1463         TestResultCallback callback = new TestResultCallback();
1464         mAppSearchManagerServiceStub.getNextPage(new GetNextPageAidlRequest(
1465                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
1466                 /* databaseName= */ null,/* nextPageToken= */ 0,
1467                 AppSearchSchema.StringPropertyConfig.JOINABLE_VALUE_TYPE_QUALIFIED_ID, mUserHandle,
1468                 BINDER_CALL_START_TIME, /* isForEnterprise= */ true), callback);
1469         AppSearchResult<SearchResultPage> result =
1470                 (AppSearchResult<SearchResultPage>) callback.get();
1471         assertThat(result.getResultCode()).isEqualTo(AppSearchResult.RESULT_OK);
1472         assertThat(result.getResultValue().getResults()).isEmpty();
1473         // No CallStats logged since we returned early
1474         verify(mLogger, timeout(1000).times(0)).logStats(any(CallStats.class));
1475     }
1476 
1477     @Test
testEnterpriseInvalidateNextPageToken_noEnterpriseUser()1478     public void testEnterpriseInvalidateNextPageToken_noEnterpriseUser() throws Exception {
1479         // Even on devices with an enterprise user, this test will run properly, since we haven't
1480         // unlocked the enterprise user for our local instance of AppSearchManagerService
1481         mAppSearchManagerServiceStub.invalidateNextPageToken(new InvalidateNextPageTokenAidlRequest(
1482                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
1483                 /* nextPageToken= */ 0, mUserHandle, BINDER_CALL_START_TIME,
1484                 /* isForEnterprise= */ true));
1485         // No CallStats logged since we returned early
1486         verify(mLogger, timeout(1000).times(0)).logStats(any(CallStats.class));
1487     }
1488 
1489     @Test
executeAppFunction_success()1490     public void executeAppFunction_success() throws Exception {
1491         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_OK);
1492     }
1493 
1494     @Test
executeAppFunction_callerNoPermission()1495     public void executeAppFunction_callerNoPermission() throws Exception {
1496         doReturn(List.of())
1497                 .when(mRoleManager).getRoleHoldersAsUser(
1498                         AppSearchManagerService.SYSTEM_UI_INTELLIGENCE, mUserHandle);
1499 
1500         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_SECURITY_ERROR);
1501     }
1502 
1503     @Test
executeAppFunction_cannotResolveService()1504     public void executeAppFunction_cannotResolveService() throws Exception {
1505         PackageManager spyPackageManager = mContext.getPackageManager();
1506         doReturn(null).when(spyPackageManager).resolveService(any(Intent.class), eq(0));
1507 
1508         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_NOT_FOUND);
1509     }
1510 
1511     @Test
executeAppFunction_serviceNotPermissionProtected()1512     public void executeAppFunction_serviceNotPermissionProtected() throws Exception {
1513         ServiceInfo serviceInfo = new ServiceInfo();
1514         serviceInfo.packageName = FOO_PACKAGE_NAME;
1515         serviceInfo.name = ".MyAppFunctionService";
1516         ResolveInfo resolveInfo = new ResolveInfo();
1517         resolveInfo.serviceInfo = serviceInfo;
1518         PackageManager spyPackageManager = mContext.getPackageManager();
1519         doReturn(resolveInfo).when(spyPackageManager).resolveService(any(Intent.class), eq(0));
1520 
1521         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_NOT_FOUND);
1522     }
1523 
1524     @Test
executeAppFunction_bindServiceReturnsFalse()1525     public void executeAppFunction_bindServiceReturnsFalse() throws Exception {
1526         mServiceCallHelper.setBindServiceResult(false);
1527         mServiceCallHelper.setOnRunServiceCallListener((callback) -> {});
1528 
1529         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_INTERNAL_ERROR);
1530     }
1531 
1532     @Test
executeAppFunction_failedToConnectService()1533     public void executeAppFunction_failedToConnectService() throws Exception {
1534         mServiceCallHelper.setOnRunServiceCallListener(
1535                 ServiceCallHelper.RunServiceCallCallback::onFailedToConnect);
1536 
1537         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_INTERNAL_ERROR);
1538     }
1539 
1540     @Test
executeAppFunction_serviceConnectionTimeout()1541     public void executeAppFunction_serviceConnectionTimeout() throws Exception {
1542         mServiceCallHelper.setOnRunServiceCallListener(
1543                 ServiceCallHelper.RunServiceCallCallback::onTimedOut);
1544 
1545         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_TIMED_OUT);
1546     }
1547 
1548     @Test
executeAppFunction_executeAppFunctionReturnsFailure()1549     public void executeAppFunction_executeAppFunctionReturnsFailure() throws Exception {
1550         mServiceCallHelper.setOnRunServiceCallListener(
1551                 (callback) -> callback.onServiceConnected(new TestableAppFunctionService(
1552                         AppSearchResult.newFailedResult(AppSearchResult.RESULT_INVALID_ARGUMENT,
1553                                 "errorMessage")), () -> {
1554                 }));
1555 
1556         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_INVALID_ARGUMENT);
1557     }
1558 
1559     @Test
executeAppFunction_hasDeviceOwner_fail()1560     public void executeAppFunction_hasDeviceOwner_fail() throws Exception {
1561         doReturn(true).when(mDevicePolicyManager).isDeviceManaged();
1562 
1563         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_SECURITY_ERROR);
1564     }
1565 
1566     @Test
executeAppFunction_fromManagedProfile_fail()1567     public void executeAppFunction_fromManagedProfile_fail() throws Exception {
1568         UserManager spyUserManager = mContext.getSystemService(UserManager.class);
1569         doReturn(true).when(spyUserManager).isManagedProfile(mUserHandle.getIdentifier());
1570 
1571         verifyExecuteAppFunctionCallbackResult(AppSearchResult.RESULT_SECURITY_ERROR);
1572     }
1573 
verifyLocalCallsResults(int resultCode)1574     private void verifyLocalCallsResults(int resultCode) throws Exception {
1575         // These APIs are local calls since they specify a database. If the API specifies a target
1576         // package, then the target package matches the calling package
1577         verifySetSchemaResult(resultCode);
1578         verifyLocalGetSchemaResult(resultCode);
1579         verifyGetNamespacesResult(resultCode);
1580         verifyPutDocumentsResult(resultCode);
1581         verifyLocalGetDocumentsResult(resultCode);
1582         verifySearchResult(resultCode);
1583         verifyLocalGetNextPageResult(resultCode);
1584         verifyWriteSearchResultsToFileResult(resultCode);
1585         verifyPutDocumentsFromFileResult(resultCode);
1586         verifySearchSuggestionResult(resultCode);
1587         verifyLocalReportUsageResult(resultCode);
1588         verifyRemoveByDocumentIdResult(resultCode);
1589         verifyRemoveByQueryResult(resultCode);
1590         verifyGetStorageInfoResult(resultCode);
1591     }
1592 
verifyGlobalCallsResults(int resultCode)1593     private void verifyGlobalCallsResults(int resultCode) throws Exception {
1594         // These APIs are global calls since either they do not specify a database or if they do,
1595         // they specify the database along with a target package that does not match the calling
1596         // package
1597         verifyGlobalGetSchemaResult(resultCode);
1598         verifyGlobalGetDocumentsResult(resultCode);
1599         verifyGlobalSearchResult(resultCode);
1600         verifyGlobalGetNextPageResult(resultCode);
1601         verifyInvalidateNextPageTokenResult(resultCode);
1602         verifyGlobalReportUsageResult(resultCode);
1603         verifyPersistToDiskResult(resultCode);
1604         verifyRegisterObserverCallbackResult(resultCode);
1605         verifyUnregisterObserverCallbackResult(resultCode);
1606         verifyInitializeResult(resultCode);
1607         verifyExecuteAppFunctionCallbackResult(resultCode);
1608     }
1609 
verifySetSchemaResult(int resultCode)1610     private void verifySetSchemaResult(int resultCode) throws Exception {
1611         TestResultCallback callback = new TestResultCallback();
1612         mAppSearchManagerServiceStub.setSchema(
1613                 new SetSchemaAidlRequest(
1614                 AppSearchAttributionSource
1615                     .createAttributionSource(mContext, mCallingPid), DATABASE_NAME,
1616                 /* schemaBundles= */ Collections.emptyList(),
1617                 /* visibilityBundles= */ Collections.emptyList(), /* forceOverride= */ false,
1618                 /* schemaVersion= */ 0, mUserHandle, BINDER_CALL_START_TIME,
1619                 SchemaMigrationStats.FIRST_CALL_GET_INCOMPATIBLE),
1620                 callback);
1621         verifyCallResult(resultCode, CallStats.CALL_TYPE_SET_SCHEMA, callback.get());
1622     }
1623 
verifyLocalGetSchemaResult(int resultCode)1624     private void verifyLocalGetSchemaResult(int resultCode) throws Exception {
1625         TestResultCallback callback = new TestResultCallback();
1626         mAppSearchManagerServiceStub.getSchema(
1627                 new GetSchemaAidlRequest(
1628                         AppSearchAttributionSource.createAttributionSource(mContext,
1629                                 mCallingPid),
1630                         mContext.getPackageName(), DATABASE_NAME, mUserHandle,
1631                         BINDER_CALL_START_TIME, /* isForEnterprise= */ false),
1632                 callback);
1633         verifyCallResult(resultCode, CallStats.CALL_TYPE_GET_SCHEMA, callback.get());
1634     }
1635 
verifyGlobalGetSchemaResult(int resultCode)1636     private void verifyGlobalGetSchemaResult(int resultCode) throws Exception {
1637         String otherPackageName = mContext.getPackageName() + "foo";
1638         TestResultCallback callback = new TestResultCallback();
1639         mAppSearchManagerServiceStub.getSchema(
1640                 new GetSchemaAidlRequest(
1641                         AppSearchAttributionSource.createAttributionSource(mContext,
1642                                 mCallingPid),
1643                         otherPackageName, DATABASE_NAME, mUserHandle, BINDER_CALL_START_TIME,
1644                         /* isForEnterprise= */ false),
1645                 callback);
1646         verifyCallResult(resultCode, CallStats.CALL_TYPE_GLOBAL_GET_SCHEMA, callback.get());
1647     }
1648 
verifyGetNamespacesResult(int resultCode)1649     private void verifyGetNamespacesResult(int resultCode) throws Exception {
1650         TestResultCallback callback = new TestResultCallback();
1651         mAppSearchManagerServiceStub.getNamespaces(
1652                 new GetNamespacesAidlRequest(
1653                         AppSearchAttributionSource.createAttributionSource(mContext,
1654                                 mCallingPid), DATABASE_NAME,
1655                         mUserHandle, BINDER_CALL_START_TIME),
1656                 callback);
1657         verifyCallResult(resultCode, CallStats.CALL_TYPE_GET_NAMESPACES, callback.get());
1658     }
1659 
verifyPutDocumentsResult(int resultCode)1660     private void verifyPutDocumentsResult(int resultCode) throws Exception {
1661         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
1662         mAppSearchManagerServiceStub.putDocuments(
1663                 new PutDocumentsAidlRequest(
1664                         AppSearchAttributionSource.createAttributionSource(mContext,
1665                                 mCallingPid), DATABASE_NAME,
1666                         new DocumentsParcel(Collections.emptyList(), Collections.emptyList()),
1667                         mUserHandle, BINDER_CALL_START_TIME), callback);
1668         verifyCallResult(resultCode, CallStats.CALL_TYPE_PUT_DOCUMENTS, callback.get());
1669     }
1670 
verifyLocalGetDocumentsResult(int resultCode)1671     private void verifyLocalGetDocumentsResult(int resultCode) throws Exception {
1672         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
1673         mAppSearchManagerServiceStub.getDocuments(
1674                 new GetDocumentsAidlRequest(
1675                         AppSearchAttributionSource.createAttributionSource(mContext,
1676                                 mCallingPid),
1677                         mContext.getPackageName(), DATABASE_NAME,
1678                         new GetByDocumentIdRequest.Builder(NAMESPACE)
1679                                 .addIds(/* ids= */ Collections.emptyList())
1680                                 .build(),
1681                         mUserHandle, BINDER_CALL_START_TIME, /* isForEnterprise= */ false),
1682                 callback);
1683         verifyCallResult(resultCode, CallStats.CALL_TYPE_GET_DOCUMENTS, callback.get());
1684     }
1685 
verifyGlobalGetDocumentsResult(int resultCode)1686     private void verifyGlobalGetDocumentsResult(int resultCode) throws Exception {
1687         String otherPackageName = mContext.getPackageName() + "foo";
1688         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
1689         mAppSearchManagerServiceStub.getDocuments(
1690                 new GetDocumentsAidlRequest(
1691                         AppSearchAttributionSource.createAttributionSource(mContext,
1692                                 mCallingPid),
1693                         otherPackageName, DATABASE_NAME,
1694                         new GetByDocumentIdRequest.Builder(NAMESPACE)
1695                                 .addIds(/* ids= */ Collections.emptyList())
1696                                 .build(),
1697                         mUserHandle, BINDER_CALL_START_TIME, /* isForEnterprise= */ false),
1698                 callback);
1699         verifyCallResult(resultCode, CallStats.CALL_TYPE_GLOBAL_GET_DOCUMENT_BY_ID, callback.get());
1700     }
1701 
verifySearchResult(int resultCode)1702     private void verifySearchResult(int resultCode) throws Exception {
1703         TestResultCallback callback = new TestResultCallback();
1704         mAppSearchManagerServiceStub.search(
1705                 new SearchAidlRequest(AppSearchAttributionSource.createAttributionSource(mContext,
1706                         mCallingPid),
1707                         DATABASE_NAME,/* searchExpression= */ "", EMPTY_SEARCH_SPEC, mUserHandle,
1708                         BINDER_CALL_START_TIME), callback);
1709         verifyCallResult(resultCode, CallStats.CALL_TYPE_SEARCH, callback.get());
1710     }
1711 
verifyGlobalSearchResult(int resultCode)1712     private void verifyGlobalSearchResult(int resultCode) throws Exception {
1713         TestResultCallback callback = new TestResultCallback();
1714         mAppSearchManagerServiceStub.globalSearch(new GlobalSearchAidlRequest(
1715                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
1716                 /* searchExpression= */ "", EMPTY_SEARCH_SPEC, mUserHandle, BINDER_CALL_START_TIME,
1717                 /* isForEnterprise= */ false), callback);
1718         verifyCallResult(resultCode, CallStats.CALL_TYPE_GLOBAL_SEARCH, callback.get());
1719     }
1720 
verifyLocalGetNextPageResult(int resultCode)1721     private void verifyLocalGetNextPageResult(int resultCode) throws Exception {
1722         TestResultCallback callback = new TestResultCallback();
1723         mAppSearchManagerServiceStub.getNextPage(new GetNextPageAidlRequest(
1724                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
1725                 DATABASE_NAME,
1726                 /* nextPageToken= */ 0,
1727                 AppSearchSchema.StringPropertyConfig.JOINABLE_VALUE_TYPE_QUALIFIED_ID, mUserHandle,
1728                 BINDER_CALL_START_TIME, /* isForEnterprise= */ false), callback);
1729         verifyCallResult(resultCode, CallStats.CALL_TYPE_GET_NEXT_PAGE, callback.get());
1730     }
1731 
verifyGlobalGetNextPageResult(int resultCode)1732     private void verifyGlobalGetNextPageResult(int resultCode) throws Exception {
1733         TestResultCallback callback = new TestResultCallback();
1734         mAppSearchManagerServiceStub.getNextPage(new GetNextPageAidlRequest(
1735                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
1736                 /* databaseName= */ null, /* nextPageToken= */ 0,
1737                 AppSearchSchema.StringPropertyConfig.JOINABLE_VALUE_TYPE_QUALIFIED_ID, mUserHandle,
1738                 BINDER_CALL_START_TIME, /* isForEnterprise= */ false), callback);
1739         verifyCallResult(resultCode, CallStats.CALL_TYPE_GLOBAL_GET_NEXT_PAGE, callback.get());
1740     }
1741 
verifyInvalidateNextPageTokenResult(int resultCode)1742     private void verifyInvalidateNextPageTokenResult(int resultCode) throws Exception {
1743         mAppSearchManagerServiceStub.invalidateNextPageToken(new InvalidateNextPageTokenAidlRequest(
1744                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
1745                 /* nextPageToken= */ 0, mUserHandle, BINDER_CALL_START_TIME,
1746                 /* isForEnterprise= */ false));
1747         verifyCallResult(resultCode, CallStats.CALL_TYPE_INVALIDATE_NEXT_PAGE_TOKEN, /* result= */
1748                 null);
1749     }
1750 
verifyWriteSearchResultsToFileResult(int resultCode)1751     private void verifyWriteSearchResultsToFileResult(int resultCode) throws Exception {
1752         File tempFile = mTemporaryFolder.newFile();
1753         FileDescriptor fd = IoBridge.open(tempFile.getPath(), O_WRONLY);
1754         TestResultCallback callback = new TestResultCallback();
1755         mAppSearchManagerServiceStub.writeSearchResultsToFile(
1756                 new WriteSearchResultsToFileAidlRequest(
1757                         AppSearchAttributionSource.createAttributionSource(mContext,
1758                                 mCallingPid), DATABASE_NAME,
1759                         new ParcelFileDescriptor(fd), /* searchExpression= */ "", EMPTY_SEARCH_SPEC,
1760                         mUserHandle, BINDER_CALL_START_TIME), callback);
1761         verifyCallResult(resultCode, CallStats.CALL_TYPE_WRITE_SEARCH_RESULTS_TO_FILE,
1762                 callback.get());
1763     }
1764 
verifyPutDocumentsFromFileResult(int resultCode)1765     private void verifyPutDocumentsFromFileResult(int resultCode) throws Exception {
1766         File tempFile = mTemporaryFolder.newFile();
1767         FileDescriptor fd = IoBridge.open(tempFile.getPath(), O_RDONLY);
1768         TestResultCallback callback = new TestResultCallback();
1769         mAppSearchManagerServiceStub.putDocumentsFromFile(new PutDocumentsFromFileAidlRequest(
1770                 AppSearchAttributionSource.createAttributionSource(mContext, mCallingPid),
1771                 DATABASE_NAME,
1772                 new ParcelFileDescriptor(fd), mUserHandle,
1773                 new SchemaMigrationStats.Builder(mContext.getPackageName(), DATABASE_NAME).build(),
1774                 /* totalLatencyStartTimeMillis= */ 0, BINDER_CALL_START_TIME), callback);
1775         verifyCallResult(resultCode, CallStats.CALL_TYPE_PUT_DOCUMENTS_FROM_FILE, callback.get());
1776     }
1777 
verifySearchSuggestionResult(int resultCode)1778     private void verifySearchSuggestionResult(int resultCode) throws Exception {
1779         SearchSuggestionSpec searchSuggestionSpec =
1780             new SearchSuggestionSpec.Builder(/*maximumResultCount=*/1).build();
1781         TestResultCallback callback = new TestResultCallback();
1782         mAppSearchManagerServiceStub.searchSuggestion(
1783                 new SearchSuggestionAidlRequest(
1784                         AppSearchAttributionSource.createAttributionSource(mContext,
1785                                 mCallingPid),
1786                         DATABASE_NAME, /* suggestionQueryExpression= */ "foo", searchSuggestionSpec,
1787                         mUserHandle, BINDER_CALL_START_TIME),
1788                 callback);
1789         verifyCallResult(resultCode, CallStats.CALL_TYPE_SEARCH_SUGGESTION, callback.get());
1790     }
1791 
verifyLocalReportUsageResult(int resultCode)1792     private void verifyLocalReportUsageResult(int resultCode) throws Exception {
1793         setUpTestSchema(mContext.getPackageName(), DATABASE_NAME);
1794         setUpTestDocument(mContext.getPackageName(), DATABASE_NAME, NAMESPACE, ID);
1795         TestResultCallback callback = new TestResultCallback();
1796         mAppSearchManagerServiceStub.reportUsage(
1797                 new ReportUsageAidlRequest(
1798                         AppSearchAttributionSource.createAttributionSource(mContext,
1799                                 mCallingPid),
1800                 mContext.getPackageName(), DATABASE_NAME,
1801                         new ReportUsageRequest.Builder(NAMESPACE, ID)
1802                                 .setUsageTimestampMillis(/* usageTimestampMillis= */ 0)
1803                                 .build(),
1804                         /* systemUsage= */ false, mUserHandle, BINDER_CALL_START_TIME),
1805                 callback);
1806         verifyCallResult(resultCode, CallStats.CALL_TYPE_REPORT_USAGE, callback.get());
1807         removeTestSchema(mContext.getPackageName(), DATABASE_NAME);
1808     }
1809 
verifyGlobalReportUsageResult(int resultCode)1810     private void verifyGlobalReportUsageResult(int resultCode) throws Exception {
1811         // Grant system access for global report usage
1812         mUiAutomation.adoptShellPermissionIdentity(Manifest.permission.READ_GLOBAL_APP_SEARCH_DATA);
1813         try {
1814             String otherPackageName = mContext.getPackageName() + "foo";
1815             setUpTestSchema(otherPackageName, DATABASE_NAME);
1816             setUpTestDocument(otherPackageName, DATABASE_NAME, NAMESPACE, ID);
1817             TestResultCallback callback = new TestResultCallback();
1818             mAppSearchManagerServiceStub.reportUsage(
1819                     new ReportUsageAidlRequest(
1820                             AppSearchAttributionSource.createAttributionSource(mContext,
1821                                     mCallingPid),
1822                             otherPackageName, DATABASE_NAME,
1823                             new ReportUsageRequest.Builder(NAMESPACE, ID)
1824                                     .setUsageTimestampMillis(/* usageTimestampMillis= */ 0)
1825                                     .build(),
1826                     /* systemUsage= */ true, mUserHandle, BINDER_CALL_START_TIME),
1827                     callback);
1828             verifyCallResult(resultCode, CallStats.CALL_TYPE_REPORT_SYSTEM_USAGE, callback.get());
1829             removeTestSchema(otherPackageName, DATABASE_NAME);
1830         } finally {
1831             mUiAutomation.dropShellPermissionIdentity();
1832         }
1833     }
1834 
verifyRemoveByDocumentIdResult(int resultCode)1835     private void verifyRemoveByDocumentIdResult(int resultCode) throws Exception {
1836         TestBatchResultErrorCallback callback = new TestBatchResultErrorCallback();
1837         mAppSearchManagerServiceStub.removeByDocumentId(
1838                 new RemoveByDocumentIdAidlRequest(
1839                         AppSearchAttributionSource.createAttributionSource(mContext,
1840                                 mCallingPid),
1841                         DATABASE_NAME,
1842                         new RemoveByDocumentIdRequest.Builder(NAMESPACE)
1843                                 .addIds(/* ids= */ Collections.emptyList())
1844                                 .build(),
1845                         mUserHandle,
1846                         BINDER_CALL_START_TIME),
1847                 callback);
1848         verifyCallResult(resultCode, CallStats.CALL_TYPE_REMOVE_DOCUMENTS_BY_ID, callback.get());
1849     }
1850 
verifyRemoveByQueryResult(int resultCode)1851     private void verifyRemoveByQueryResult(int resultCode) throws Exception {
1852         TestResultCallback callback = new TestResultCallback();
1853         mAppSearchManagerServiceStub.removeByQuery(
1854                 new RemoveByQueryAidlRequest(
1855                         AppSearchAttributionSource.createAttributionSource(mContext,
1856                                 mCallingPid), DATABASE_NAME,
1857                         /* queryExpression= */ "", EMPTY_SEARCH_SPEC, mUserHandle,
1858                         BINDER_CALL_START_TIME),
1859                 callback);
1860         verifyCallResult(resultCode, CallStats.CALL_TYPE_REMOVE_DOCUMENTS_BY_SEARCH,
1861                 callback.get());
1862     }
1863 
verifyGetStorageInfoResult(int resultCode)1864     private void verifyGetStorageInfoResult(int resultCode) throws Exception {
1865         TestResultCallback callback = new TestResultCallback();
1866         mAppSearchManagerServiceStub.getStorageInfo(
1867                 new GetStorageInfoAidlRequest(
1868                         AppSearchAttributionSource.createAttributionSource(mContext,
1869                                 mCallingPid), DATABASE_NAME,
1870                         mUserHandle, BINDER_CALL_START_TIME),
1871                 callback);
1872         verifyCallResult(resultCode, CallStats.CALL_TYPE_GET_STORAGE_INFO, callback.get());
1873     }
1874 
verifyPersistToDiskResult(int resultCode)1875     private void verifyPersistToDiskResult(int resultCode) throws Exception {
1876         mAppSearchManagerServiceStub.persistToDisk(
1877                 new PersistToDiskAidlRequest(
1878                         AppSearchAttributionSource.createAttributionSource(mContext,
1879                                 mCallingPid), mUserHandle,
1880                         BINDER_CALL_START_TIME));
1881         verifyCallResult(resultCode, CallStats.CALL_TYPE_FLUSH, /* result= */ null);
1882     }
1883 
verifyRegisterObserverCallbackResult(int resultCode)1884     private void verifyRegisterObserverCallbackResult(int resultCode) throws Exception {
1885         AppSearchResultParcel<Void> resultParcel =
1886                 mAppSearchManagerServiceStub.registerObserverCallback(
1887                         new RegisterObserverCallbackAidlRequest(
1888                                 AppSearchAttributionSource.createAttributionSource(mContext,
1889                                         mCallingPid),
1890                                 mContext.getPackageName(),
1891                                 new ObserverSpec.Builder().build(),
1892                                 mUserHandle,
1893                                 BINDER_CALL_START_TIME),
1894                         new IAppSearchObserverProxy.Stub() {
1895                             @Override
1896                             public void onSchemaChanged(String packageName, String databaseName,
1897                                     List<String> changedSchemaNames) throws RemoteException {
1898                             }
1899 
1900                             @Override
1901                             public void onDocumentChanged(String packageName, String databaseName,
1902                                     String namespace, String schemaName,
1903                                     List<String> changedDocumentIds) throws RemoteException {
1904                             }
1905                         });
1906         verifyCallResult(resultCode, CallStats.CALL_TYPE_REGISTER_OBSERVER_CALLBACK,
1907                 resultParcel.getResult());
1908     }
1909 
verifyExecuteAppFunctionCallbackResult(int resultCode)1910     private void verifyExecuteAppFunctionCallbackResult(int resultCode) throws Exception {
1911         TestResultCallback callback = new TestResultCallback();
1912         mAppSearchManagerServiceStub.executeAppFunction(
1913                 new ExecuteAppFunctionAidlRequest(
1914                         new ExecuteAppFunctionRequest.Builder(
1915                                 FOO_PACKAGE_NAME, "function"
1916                         ).build(),
1917                         AppSearchAttributionSource.createAttributionSource(mContext,
1918                                 mCallingPid),
1919                         mUserHandle,
1920                         BINDER_CALL_START_TIME
1921                 ),
1922                 callback
1923         );
1924 
1925         verifyCallResult(resultCode, CallStats.CALL_TYPE_EXECUTE_APP_FUNCTION, callback.get());
1926     }
1927 
verifyUnregisterObserverCallbackResult(int resultCode)1928     private void verifyUnregisterObserverCallbackResult(int resultCode) throws Exception {
1929         AppSearchResultParcel<Void> resultParcel =
1930                 mAppSearchManagerServiceStub.unregisterObserverCallback(
1931                         new UnregisterObserverCallbackAidlRequest(
1932                                 AppSearchAttributionSource.createAttributionSource(mContext,
1933                                         mCallingPid),
1934                                 mContext.getPackageName(), mUserHandle,
1935                                 BINDER_CALL_START_TIME),
1936                         new IAppSearchObserverProxy.Stub() {
1937                             @Override
1938                             public void onSchemaChanged(String packageName, String databaseName,
1939                                     List<String> changedSchemaNames) throws RemoteException {
1940                             }
1941 
1942                             @Override
1943                             public void onDocumentChanged(String packageName, String databaseName,
1944                                     String namespace, String schemaName,
1945                                     List<String> changedDocumentIds) throws RemoteException {
1946                             }
1947                         });
1948         verifyCallResult(resultCode, CallStats.CALL_TYPE_UNREGISTER_OBSERVER_CALLBACK,
1949                 resultParcel.getResult());
1950     }
1951 
verifyInitializeResult(int resultCode)1952     private void verifyInitializeResult(int resultCode) throws Exception {
1953         TestResultCallback callback = new TestResultCallback();
1954         mAppSearchManagerServiceStub.initialize(
1955                 new InitializeAidlRequest(
1956                         AppSearchAttributionSource.createAttributionSource(mContext,
1957                                 mCallingPid), mUserHandle,
1958                         BINDER_CALL_START_TIME),
1959                 callback);
1960         if (resultCode == RESULT_DENIED) {
1961             verify(mLogger, never()).logStats(any(CallStats.class));
1962         } else {
1963             verifyCallResult(resultCode, CallStats.CALL_TYPE_INITIALIZE, /* result= */ null);
1964         }
1965         assertThat(callback.get().getResultCode()).isEqualTo(resultCode);
1966     }
1967 
verifyCallResult(int resultCode, int callType, AppSearchResult<?> result)1968     private void verifyCallResult(int resultCode, int callType, AppSearchResult<?> result) {
1969         ArgumentCaptor<CallStats> captor = ArgumentCaptor.forClass(CallStats.class);
1970         verify(mLogger, timeout(1000).times(1)).logStats(captor.capture());
1971         assertThat(captor.getValue().getCallType()).isEqualTo(callType);
1972         assertThat(captor.getValue().getStatusCode()).isEqualTo(resultCode);
1973         assertThat(captor.getValue().getEstimatedBinderLatencyMillis()).isGreaterThan(0);
1974         clearInvocations(mLogger);
1975         // Not all calls return a result
1976         if (result != null) {
1977             assertThat(result.getResultCode()).isEqualTo(resultCode);
1978         }
1979     }
1980 
setUpTestSchema(String packageName, String databaseName)1981     private void setUpTestSchema(String packageName, String databaseName) throws Exception {
1982         // Insert schema
1983         List<AppSearchSchema> schemas = Collections.singletonList(
1984                 new AppSearchSchema.Builder("type").build());
1985         InternalSetSchemaResponse internalSetSchemaResponse =
1986                 mUserInstance.getAppSearchImpl().setSchema(packageName, databaseName, schemas,
1987                         /* visibilityDocuments= */ Collections.emptyList(),
1988                         /* forceOverride= */ false,
1989                         /* version= */ 0,
1990                         /* setSchemaStatsBuilder= */ null);
1991         assertThat(internalSetSchemaResponse.isSuccess()).isTrue();
1992     }
1993 
setUpTestDocument(String packageName, String databaseName, String namespace, String id)1994     private void setUpTestDocument(String packageName, String databaseName, String namespace,
1995             String id) throws Exception {
1996         // Insert a document
1997         GenericDocument document = new GenericDocument.Builder<>(namespace, id, "type").build();
1998         mUserInstance.getAppSearchImpl().putDocument(packageName, databaseName, document,
1999                 /* sendChangeNotifications= */ false,
2000                 /* logger= */ null);
2001     }
2002 
removeTestSchema(String packageName, String databaseName)2003     private void removeTestSchema(String packageName, String databaseName) throws Exception {
2004         InternalSetSchemaResponse internalSetSchemaResponse =
2005                 mUserInstance.getAppSearchImpl().setSchema(packageName, databaseName,
2006                         /* schemas= */ Collections.emptyList(),
2007                         /* visibilityDocuments= */ Collections.emptyList(),
2008                         /* forceOverride= */ true,
2009                         /* version= */ 0,
2010                         /* setSchemaStatsBuilder= */ null);
2011         assertThat(internalSetSchemaResponse.isSuccess()).isTrue();
2012     }
2013 
setUpEnvironmentForAppFunction()2014     private void setUpEnvironmentForAppFunction() {
2015         doReturn(Arrays.asList(mContext.getPackageName()))
2016                 .when(mRoleManager).getRoleHoldersAsUser(
2017                         AppSearchManagerService.SYSTEM_UI_INTELLIGENCE, mUserHandle);
2018 
2019         // FOO_PACKAGE implemented an AppFunctionService.
2020         ServiceInfo serviceInfo = new ServiceInfo();
2021         serviceInfo.packageName = FOO_PACKAGE_NAME;
2022         serviceInfo.name = ".MyAppFunctionService";
2023         serviceInfo.permission = AppFunctionManager.PERMISSION_BIND_APP_FUNCTION_SERVICE;
2024         ResolveInfo resolveInfo = new ResolveInfo();
2025         resolveInfo.serviceInfo = serviceInfo;
2026         PackageManager spyPackageManager = mContext.getPackageManager();
2027         doReturn(resolveInfo).when(spyPackageManager).resolveService(any(Intent.class), eq(0));
2028 
2029         doReturn(false).when(mDevicePolicyManager).isDeviceManaged();
2030         UserManager spyUserManager = mContext.getSystemService(UserManager.class);
2031         doReturn(false).when(spyUserManager).isManagedProfile(mUserHandle.getIdentifier());
2032     }
2033 
2034     private static class MockServiceManager implements StaticMockFixture {
2035         ArgumentCaptor<IAppSearchManager.Stub> mStubCaptor = ArgumentCaptor.forClass(
2036                 IAppSearchManager.Stub.class);
2037 
2038         @Override
setUpMockedClasses( @onNull StaticMockitoSessionBuilder sessionBuilder)2039         public StaticMockitoSessionBuilder setUpMockedClasses(
2040                 @NonNull StaticMockitoSessionBuilder sessionBuilder) {
2041             sessionBuilder.mockStatic(LocalManagerRegistry.class);
2042             sessionBuilder.spyStatic(ServiceManager.class);
2043             return sessionBuilder;
2044         }
2045 
2046         @Override
setUpMockBehaviors()2047         public void setUpMockBehaviors() {
2048             ExtendedMockito.doReturn(mock(StorageStatsManagerLocal.class)).when(
2049                     () -> LocalManagerRegistry.getManager(StorageStatsManagerLocal.class));
2050             ExtendedMockito.doNothing().when(
2051                     () -> ServiceManager.addService(anyString(), mStubCaptor.capture(),
2052                             anyBoolean(), anyInt()));
2053         }
2054 
2055         @Override
tearDown()2056         public void tearDown() {
2057         }
2058     }
2059 
2060     private static final class TestResultCallback extends IAppSearchResultCallback.Stub {
2061         private final SettableFuture<AppSearchResult<?>> future = SettableFuture.create();
2062 
2063         @Override
onResult(AppSearchResultParcel appSearchResultParcel)2064         public void onResult(AppSearchResultParcel appSearchResultParcel) {
2065             future.set(appSearchResultParcel.getResult());
2066         }
2067 
get()2068         public AppSearchResult<?> get() throws InterruptedException, ExecutionException {
2069             return future.get();
2070         }
2071     }
2072 
2073     private static final class TestBatchResultErrorCallback extends
2074             IAppSearchBatchResultCallback.Stub {
2075         private final SettableFuture<AppSearchResult<?>> future = SettableFuture.create();
2076         private final SettableFuture<AppSearchBatchResult<?, ?>> batchFuture =
2077                 SettableFuture.create();
2078 
2079         @Override
onResult(AppSearchBatchResultParcel appSearchBatchResultParcel)2080         public void onResult(AppSearchBatchResultParcel appSearchBatchResultParcel) {
2081             future.set(null);
2082             batchFuture.set(appSearchBatchResultParcel.getResult());
2083         }
2084 
2085         @Override
onSystemError(AppSearchResultParcel appSearchResultParcel)2086         public void onSystemError(AppSearchResultParcel appSearchResultParcel) {
2087             future.set(appSearchResultParcel.getResult());
2088             batchFuture.set(null);
2089         }
2090 
get()2091         public AppSearchResult<?> get() throws InterruptedException, ExecutionException {
2092             return future.get();
2093         }
2094 
getBatchResult()2095         public AppSearchBatchResult<?, ?> getBatchResult()
2096                 throws InterruptedException, ExecutionException {
2097             return batchFuture.get();
2098         }
2099     }
2100 
2101     /**
2102      * Testable helper for {@link ServiceCallHelper}, defaults to the happy case, i.e. successful
2103      * service connection and
2104      * {@link android.app.appsearch.functions.AppFunctionService#onExecuteFunction} always
2105      * returns a successful result. Includes methods to customize connection behavior.
2106      */
2107     private static class TestableServiceCallHelper implements
2108             ServiceCallHelper<IAppFunctionService> {
2109         private Consumer<RunServiceCallCallback<IAppFunctionService>> mOnRunServiceCallListener =
2110                 (callback) -> callback.onServiceConnected(new TestableAppFunctionService(
2111                         AppSearchResult.newSuccessfulResult(
2112                                 new ExecuteAppFunctionResponse.Builder().build())), () -> {
2113                 });
2114         private boolean mBindServiceResult = true;
2115 
2116         /**
2117          * Replaces the default service connection behavior. Use this in tests to simulate
2118          * different connection results (e.g., failures).
2119          */
setOnRunServiceCallListener( @onNull Consumer<RunServiceCallCallback<IAppFunctionService>> listener)2120         public void setOnRunServiceCallListener(
2121                 @NonNull Consumer<RunServiceCallCallback<IAppFunctionService>> listener) {
2122             mOnRunServiceCallListener = Objects.requireNonNull(listener);
2123         }
2124 
2125         /** Sets the result of {@link #runServiceCall} (defaults to {@code true}). */
setBindServiceResult(boolean bindResult)2126         public void setBindServiceResult(boolean bindResult) {
2127             mBindServiceResult = bindResult;
2128         }
2129 
2130         @Override
runServiceCall(@onNull Intent intent, int bindFlags, long timeoutInMillis, @NonNull UserHandle userHandle, @NonNull RunServiceCallCallback<IAppFunctionService> callback)2131         public boolean runServiceCall(@NonNull Intent intent, int bindFlags,
2132                 long timeoutInMillis, @NonNull UserHandle userHandle,
2133                 @NonNull RunServiceCallCallback<IAppFunctionService> callback) {
2134             mOnRunServiceCallListener.accept(callback);
2135             return mBindServiceResult;
2136         }
2137     }
2138 
2139     /**
2140      * A testable implementation of {@link IAppFunctionService.Stub} for you to customize the
2141      * result of {@link #executeAppFunction}.
2142      */
2143     private static class TestableAppFunctionService extends IAppFunctionService.Stub {
2144         private final AppSearchResult<ExecuteAppFunctionResponse> mResult;
2145 
2146         /**
2147          * @param result the result to return in {@link #executeAppFunction}.
2148          */
TestableAppFunctionService( @onNull AppSearchResult<ExecuteAppFunctionResponse> result)2149         public TestableAppFunctionService(
2150                 @NonNull AppSearchResult<ExecuteAppFunctionResponse> result) {
2151             mResult = Objects.requireNonNull(result);
2152         }
2153 
2154         @Override
executeAppFunction( ExecuteAppFunctionRequest executeAppFunctionRequest, IAppSearchResultCallback callback)2155         public void executeAppFunction(
2156                 ExecuteAppFunctionRequest executeAppFunctionRequest,
2157                 IAppSearchResultCallback callback) throws RemoteException {
2158             if (mResult.isSuccess()) {
2159                 callback.onResult(AppSearchResultParcel.fromExecuteAppFunctionResponse(
2160                         mResult.getResultValue()));
2161             } else {
2162                 callback.onResult(AppSearchResultParcel.fromFailedResult(mResult));
2163             }
2164         }
2165     }
2166 }
2167