• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /*
2  * Copyright (C) 2024 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package android.ondeviceintelligence.cts;
18 
19 import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE;
20 import static android.app.ondeviceintelligence.flags.Flags.FLAG_ENABLE_ON_DEVICE_INTELLIGENCE_MODULE;
21 import static android.content.Context.RECEIVER_EXPORTED;
22 import static android.ondeviceintelligence.cts.CtsIsolatedInferenceService.constructException;
23 import static android.ondeviceintelligence.cts.CtsIsolatedInferenceService.constructTokenInfo;
24 
25 import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;
26 
27 import static com.android.compatibility.common.util.ShellUtils.runShellCommand;
28 
29 import static com.google.common.truth.Truth.assertThat;
30 
31 import static org.junit.Assert.assertEquals;
32 import static org.junit.Assert.assertThrows;
33 import static org.junit.Assume.assumeFalse;
34 import static org.junit.Assume.assumeTrue;
35 
36 import static java.util.concurrent.TimeUnit.HOURS;
37 import static java.util.concurrent.TimeUnit.SECONDS;
38 
39 import android.Manifest;
40 import android.app.ondeviceintelligence.DownloadCallback;
41 import android.app.ondeviceintelligence.Feature;
42 import android.app.ondeviceintelligence.InferenceInfo;
43 import android.app.ondeviceintelligence.OnDeviceIntelligenceException;
44 import android.app.ondeviceintelligence.OnDeviceIntelligenceManager;
45 import android.app.ondeviceintelligence.ProcessingCallback;
46 import android.app.ondeviceintelligence.ProcessingSignal;
47 import android.app.ondeviceintelligence.StreamingProcessingCallback;
48 import android.app.ondeviceintelligence.TokenInfo;
49 import android.content.BroadcastReceiver;
50 import android.content.ComponentName;
51 import android.content.Context;
52 import android.content.Intent;
53 import android.content.IntentFilter;
54 import android.content.ServiceConnection;
55 import android.content.pm.PackageManager;
56 import android.os.Bundle;
57 import android.os.CancellationSignal;
58 import android.os.IBinder;
59 import android.os.PersistableBundle;
60 import android.os.Process;
61 import android.os.UserHandle;
62 import android.platform.test.annotations.AppModeFull;
63 import android.platform.test.annotations.RequiresFlagsEnabled;
64 import android.platform.test.flag.junit.CheckFlagsRule;
65 import android.platform.test.flag.junit.DeviceFlagsValueProvider;
66 import android.provider.DeviceConfig;
67 import android.provider.Settings;
68 import android.service.ondeviceintelligence.OnDeviceIntelligenceService;
69 import android.text.TextUtils;
70 import android.util.Log;
71 
72 import androidx.annotation.NonNull;
73 import androidx.annotation.Nullable;
74 import androidx.test.InstrumentationRegistry;
75 import androidx.test.ext.junit.runners.AndroidJUnit4;
76 
77 import com.android.compatibility.common.util.DeviceConfigStateChangerRule;
78 
79 import org.junit.After;
80 import org.junit.Before;
81 import org.junit.Rule;
82 import org.junit.Test;
83 import org.junit.rules.TestRule;
84 import org.junit.runner.RunWith;
85 import org.junit.runners.model.Statement;
86 
87 import java.lang.annotation.ElementType;
88 import java.lang.annotation.Retention;
89 import java.lang.annotation.RetentionPolicy;
90 import java.lang.annotation.Target;
91 import java.util.List;
92 import java.util.concurrent.CompletableFuture;
93 import java.util.concurrent.CountDownLatch;
94 import java.util.concurrent.ExecutionException;
95 import java.util.concurrent.Executor;
96 import java.util.concurrent.Executors;
97 import java.util.function.Consumer;
98 
99 /**
100  * Test the OnDeviceIntelligenceManager API. Run with "atest OnDeviceIntelligenceManagerTest"
101  * .
102  */
103 @RunWith(AndroidJUnit4.class)
104 @AppModeFull(reason = "PM will not recognize OnDeviceIntelligenceManagerService in instantMode.")
105 public class OnDeviceIntelligenceManagerTest {
106     public static final String TEST_FILE_NAME = "test_file.txt";
107     public static final String TEST_KEY = "test_key";
108     public static final String TEST_CONTENT = "test_content";
109     public static final String TEST_AUGMENT_KEY = "test_augment_key";
110     public static final String TEST_AUGMENT_CONTENT = "test_augment_content";
111     public static final String EXCEPTION_MESSAGE_KEY = "message_key";
112     public static final String EXCEPTION_STATUS_CODE_KEY = "code_key";
113     public static final String EXCEPTION_PARAMS_KEY = "params_key";
114     public static final String TOKEN_INFO_COUNT_KEY = "tokenInfo_count_key";
115     public static final String TOKEN_INFO_PARAMS_KEY = "tokenInfo_params_key";
116     public static final String TEST_OD_NAMESPACE = "test_od_namespace";
117 
118 
119     private static final String TAG = OnDeviceIntelligenceManagerTest.class.getSimpleName();
120     public static final String CTS_PACKAGE_NAME =
121             android.ondeviceintelligence.cts.CtsIntelligenceService.class.getPackageName();
122     public static final String CTS_INTELLIGENCE_SERVICE_NAME =
123             CTS_PACKAGE_NAME + "/"
124                     + android.ondeviceintelligence.cts.CtsIntelligenceService.class.getCanonicalName();
125     public static final String CTS_INFERENCE_SERVICE_NAME =
126             CTS_PACKAGE_NAME + "/"
127                     + android.ondeviceintelligence.cts.CtsIsolatedInferenceService.class.getCanonicalName();
128     private static final int TEMPORARY_SERVICE_DURATION = 20000;
129     public static final String NAMESPACE_ON_DEVICE_INTELLIGENCE = "ondeviceintelligence";
130     public static final String KEY_SERVICE_ENABLED = "service_enabled";
131 
132     public static final int REQUEST_TYPE_GET_PACKAGE_NAME = 1000;
133 
134     public static final int REQUEST_TYPE_GET_FILE_FROM_MAP = 1001;
135     public static final int REQUEST_TYPE_GET_FILE_FROM_STREAM = 1002;
136     public static final int REQUEST_TYPE_GET_FILE_FROM_PFD = 1003;
137     public static final int REQUEST_TYPE_GET_AUGMENTED_DATA = 1004;
138     public static final int REQUEST_TYPE_GET_CALLER_UID = 1005;
139     public static final int REQUEST_TYPE_GET_UPDATED_DEVICE_CONFIG = 1006;
140     public static final int REQUEST_TYPE_GET_FILE_FROM_NON_FILES_DIRECTORY = 1007;
141     public static final int REQUEST_TYPE_POPULATE_INFERENCE_INFO = 1008;
142 
143     private static final Executor EXECUTOR = InstrumentationRegistry.getContext().getMainExecutor();
144     private static final String MODEL_LOADED_BROADCAST_ACTION =
145             "android.service.ondeviceintelligence.MODEL_LOADED";
146 
147     private Context mContext;
148     public OnDeviceIntelligenceManager mOnDeviceIntelligenceManager;
149     private final Executor mExecutor = Executors.newCachedThreadPool();
150 
151     @Rule
152     public final DeviceConfigStateChangerRule mDeviceConfigStateChangerRule =
153             new DeviceConfigStateChangerRule(
154                     getInstrumentation().getTargetContext(),
155                     NAMESPACE_ON_DEVICE_INTELLIGENCE,
156                     KEY_SERVICE_ENABLED,
157                     "true");
158 
159     @Rule
160     public final CheckFlagsRule mCheckFlagsRule = DeviceFlagsValueProvider.createCheckFlagsRule();
161 
162     @Before
setUp()163     public void setUp() throws Exception {
164         mContext = getInstrumentation().getContext();
165         mOnDeviceIntelligenceManager = mContext.getSystemService(OnDeviceIntelligenceManager.class);
166         bindToTestableOnDeviceIntelligenceServices();
167         setTestableDeviceConfigNamespace(TEST_OD_NAMESPACE);
168     }
169 
170     @After
tearDown()171     public void tearDown() throws Exception {
172         getInstrumentation().getUiAutomation().dropShellPermissionIdentity();
173     }
174 
175     @Test
176     @SkipSetupAndTeardown
177     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
cannotBindToIsolatedComputeAppEvenFromSamePackage()178     public void cannotBindToIsolatedComputeAppEvenFromSamePackage() {
179         assertThrows(
180                 "Cannot bind to isolated_compute_app process from same package",
181                 SecurityException.class,
182                 () -> getInstrumentation().getContext().bindService(
183                         new Intent().setComponent(new ComponentName(CTS_PACKAGE_NAME,
184                                 CtsIsolatedInferenceService.class.getCanonicalName())),
185                         new ServiceConnection() {
186                             @Override
187                             public void onServiceConnected(ComponentName name,
188                                     IBinder service) {
189                                 Log.i(TAG, "Service connected");
190                             }
191 
192                             @Override
193                             public void onServiceDisconnected(ComponentName name) {
194                                 Log.i(TAG, "Service disconnected");
195                             }
196                         },
197                         Context.BIND_AUTO_CREATE));
198     }
199 
200 //=====================Tests for Access Denied without Permission on all Manager Methods=========
201 
202     @Test
203     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
noAccessWhenAttemptingGetFeature()204     public void noAccessWhenAttemptingGetFeature() {
205         assertEquals(PackageManager.PERMISSION_DENIED, mContext.checkCallingOrSelfPermission(
206                 Manifest.permission.USE_ON_DEVICE_INTELLIGENCE));
207 
208         // Test non system app throws SecurityException
209         assertThrows("no access to getFeature from non system component",
210                 SecurityException.class,
211                 () -> mOnDeviceIntelligenceManager.getFeature(1, EXECUTOR,
212                         result -> {
213                             Log.i(TAG, "Feature : =" + result);
214                         }));
215     }
216 
217     @Test
218     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
noAccessWhenAttemptingGetFeatureDetails()219     public void noAccessWhenAttemptingGetFeatureDetails() {
220         assertEquals(PackageManager.PERMISSION_DENIED, mContext.checkCallingOrSelfPermission(
221                 Manifest.permission.USE_ON_DEVICE_INTELLIGENCE));
222         Feature feature = new Feature.Builder(1).build();
223 
224         // Test non system app throws SecurityException
225         assertThrows("no access to getFeature from non system component",
226                 SecurityException.class,
227                 () -> mOnDeviceIntelligenceManager.getFeatureDetails(feature,
228                         EXECUTOR,
229                         result -> Log.i(TAG, "Feature details : =" + result)));
230     }
231 
232     @Test
233     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
noAccessWhenAttemptingGetVersion()234     public void noAccessWhenAttemptingGetVersion() {
235         assertEquals(
236                 PackageManager.PERMISSION_DENIED,
237                 mContext.checkCallingOrSelfPermission(
238                         Manifest.permission.USE_ON_DEVICE_INTELLIGENCE));
239 
240         // Test non system app throws SecurityException
241         assertThrows(
242                 "no access to getVersion from non system component",
243                 SecurityException.class,
244                 () ->
245                         mOnDeviceIntelligenceManager.getVersion(EXECUTOR,
246                                 result -> {
247                                     Log.i(TAG, "Version : =" + result);
248                                 }));
249     }
250 
251     @Test
252     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
noAccessWhenAttemptingRequestFeatureDownload()253     public void noAccessWhenAttemptingRequestFeatureDownload() {
254         assertEquals(
255                 PackageManager.PERMISSION_DENIED,
256                 mContext.checkCallingOrSelfPermission(
257                         Manifest.permission.USE_ON_DEVICE_INTELLIGENCE));
258 
259         Feature feature = new Feature.Builder(1).build();
260 
261         // Test non system app throws SecurityException
262         assertThrows(
263                 "no access to requestFeatureDownload from non system component",
264                 SecurityException.class,
265                 () ->
266                         mOnDeviceIntelligenceManager.requestFeatureDownload(feature, null, EXECUTOR,
267                                 new DownloadCallback() {
268                                     @Override
269                                     public void onDownloadFailed(int failureStatus,
270                                             @Nullable String errorMessage,
271                                             @NonNull PersistableBundle errorParams) {
272                                         Log.e(TAG, "Got Error", new RuntimeException(errorMessage));
273                                     }
274 
275                                     @Override
276                                     public void onDownloadCompleted(
277                                             @NonNull PersistableBundle downloadParams) {
278                                         Log.i(TAG, "Response : =" + downloadParams.toString());
279                                     }
280                                 }));
281     }
282 
283     @Test
284     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
noAccessWhenRequestTokenInfo()285     public void noAccessWhenRequestTokenInfo() {
286         assertEquals(
287                 PackageManager.PERMISSION_DENIED,
288                 mContext.checkCallingOrSelfPermission(
289                         Manifest.permission.USE_ON_DEVICE_INTELLIGENCE));
290 
291 
292         Feature feature = new Feature.Builder(1).build();
293         // Test non system app throws SecurityException
294         assertThrows(
295                 "no access to requestTokenInfo from non system component",
296                 SecurityException.class,
297                 () ->
298                         mOnDeviceIntelligenceManager.requestTokenInfo(feature,
299                                 new Bundle(), null,
300                                 EXECUTOR,
301                                 result -> {
302                                     Log.i(TAG, "Response : =" + result.getCount());
303                                 }));
304     }
305 
306     @Test
307     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
noAccessWhenAttemptingProcessRequest()308     public void noAccessWhenAttemptingProcessRequest() {
309         assertEquals(
310                 PackageManager.PERMISSION_DENIED,
311                 mContext.checkCallingOrSelfPermission(
312                         Manifest.permission.USE_ON_DEVICE_INTELLIGENCE));
313 
314         Feature feature = new Feature.Builder(1).build();
315         // Test non system app throws SecurityException
316         assertThrows(
317                 "no access to processRequest from non system component",
318                 SecurityException.class,
319                 () -> mOnDeviceIntelligenceManager.processRequest(feature,
320                         new Bundle(), 1, null,
321                         null, EXECUTOR, new ProcessingCallback() {
322                             @Override
323                             public void onResult(@NonNull Bundle result) {
324                                 Log.i(TAG, "Final Result : " + result);
325                             }
326 
327                             @Override
328                             public void onError(@NonNull OnDeviceIntelligenceException error) {
329                                 Log.e(TAG, "Error Occurred", error);
330                             }
331                         }));
332     }
333 
334     @Test
335     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
noAccessWhenAttemptingProcessRequestStreaming()336     public void noAccessWhenAttemptingProcessRequestStreaming() {
337         assertEquals(
338                 PackageManager.PERMISSION_DENIED,
339                 mContext.checkCallingOrSelfPermission(
340                         Manifest.permission.USE_ON_DEVICE_INTELLIGENCE));
341 
342         Feature feature = new Feature.Builder(1).build();
343         // Test non system app throws SecurityException
344         assertThrows(
345                 "no access to processRequestStreaming from non system component",
346                 SecurityException.class,
347                 () -> mOnDeviceIntelligenceManager.processRequestStreaming(feature,
348                         new Bundle(), 1,
349                         null, null, EXECUTOR,
350                         new StreamingProcessingCallback() {
351                             @Override
352                             public void onPartialResult(@NonNull Bundle partialResult) {
353                                 Log.i(TAG, "New Content : " + partialResult);
354                             }
355 
356                             @Override
357                             public void onResult(Bundle result) {
358                                 Log.i(TAG, "Final Result : " + result);
359                             }
360 
361                             @Override
362                             public void onError(@NonNull OnDeviceIntelligenceException error) {
363                                 Log.e(TAG, "Final Result : ", error);
364                             }
365                         }));
366     }
367 
368     @Test
noAccessWhenAttemptingGetLatestInferenceInfo()369     public void noAccessWhenAttemptingGetLatestInferenceInfo() {
370         assertEquals(
371                 PackageManager.PERMISSION_DENIED,
372                 mContext.checkCallingOrSelfPermission(
373                         Manifest.permission.DUMP));
374 
375         Feature feature = new Feature.Builder(1).build();
376         // Test non system app throws SecurityException
377         assertThrows(
378                 "no access to getLatestInferenceInfo when missing permission.",
379                 SecurityException.class,
380                 () -> mOnDeviceIntelligenceManager.getLatestInferenceInfo(0));
381     }
382 
383 //===================== Tests for Result callback invoked on all Manager Methods ==================
384 
385     @Test
386     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
resultPopulatedWhenAttemptingGetFeature()387     public void resultPopulatedWhenAttemptingGetFeature() throws Exception {
388         getInstrumentation()
389                 .getUiAutomation()
390                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
391         CountDownLatch statusLatch = new CountDownLatch(1);
392         Feature expectedFeature = CtsIntelligenceService.getSampleFeature(1);
393         mOnDeviceIntelligenceManager.getFeature(1,
394                 EXECUTOR,
395                 result -> {
396                     Log.i(TAG, "Feature : =" + result);
397                     assertEquals(result.getFeatureParams().size(),
398                             expectedFeature.getFeatureParams().size());
399                     assertEquals(result.getId(), expectedFeature.getId());
400                     assertEquals(result.getName(), expectedFeature.getName());
401                     assertEquals(result.getModelName(), expectedFeature.getModelName());
402                     assertEquals(result.getType(), expectedFeature.getType());
403                     assertEquals(result.getVariant(), expectedFeature.getVariant());
404                     statusLatch.countDown();
405                 });
406         assertThat(statusLatch.await(1, SECONDS)).isTrue();
407     }
408 
409     @Test
410     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
resultPopulatedWhenAttemptingGetFeatureDetails()411     public void resultPopulatedWhenAttemptingGetFeatureDetails() throws Exception {
412         getInstrumentation()
413                 .getUiAutomation()
414                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
415         CountDownLatch statusLatch = new CountDownLatch(1);
416 
417         // Test coverage for response with params
418         mOnDeviceIntelligenceManager.getFeatureDetails(CtsIntelligenceService.getSampleFeature(0),
419                 EXECUTOR,
420                 result -> {
421                     Log.i(TAG, "Feature details : =" + result);
422                     assertEquals(result.getFeatureStatus(), 0);
423                     assertEquals(result.getFeatureDetailParams().getInt(TEST_KEY), 1);
424                     statusLatch.countDown();
425                 });
426         assertThat(statusLatch.await(1, SECONDS)).isTrue();
427 
428         // Test coverage for response withOut params
429         mOnDeviceIntelligenceManager.getFeatureDetails(CtsIntelligenceService.getSampleFeature(1),
430                 EXECUTOR,
431                 result -> {
432                     Log.i(TAG, "Feature details : =" + result);
433                     assertEquals(result.getFeatureStatus(), 1);
434                     assertEquals(result.getFeatureDetailParams().size(), 0);
435                     statusLatch.countDown();
436                 });
437         assertThat(statusLatch.await(1, SECONDS)).isTrue();
438     }
439 
440     @Test
441     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
resultPopulatedWhenAttemptingGetVersion()442     public void resultPopulatedWhenAttemptingGetVersion() throws Exception {
443         getInstrumentation()
444                 .getUiAutomation()
445                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
446         CountDownLatch statusLatch = new CountDownLatch(1);
447 
448         mOnDeviceIntelligenceManager.getVersion(EXECUTOR,
449                 result -> {
450                     Log.i(TAG, "Version : =" + result);
451                     statusLatch.countDown();
452                 });
453         assertThat(statusLatch.await(1, SECONDS)).isTrue();
454     }
455 
456     @Test
457     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
resultPopulatedWhenAttemptingRequestFeatureDownload()458     public void resultPopulatedWhenAttemptingRequestFeatureDownload() throws Exception {
459         getInstrumentation()
460                 .getUiAutomation()
461                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
462         Feature feature = new Feature.Builder(1).build();
463         CountDownLatch statusLatch = new CountDownLatch(3);
464 
465         mOnDeviceIntelligenceManager.requestFeatureDownload(feature, null, EXECUTOR,
466                 new DownloadCallback() {
467                     @Override
468                     public void onDownloadFailed(int failureStatus,
469                             @Nullable String errorMessage,
470                             @NonNull PersistableBundle errorParams) {
471                         Log.e(TAG, "Got Error", new RuntimeException(errorMessage));
472                     }
473 
474                     @Override
475                     public void onDownloadProgress(long bytesDownloaded) {
476                         statusLatch.countDown();
477                     }
478 
479                     @Override
480                     public void onDownloadStarted(long bytesDownloaded) {
481                         statusLatch.countDown();
482                     }
483 
484                     @Override
485                     public void onDownloadCompleted(
486                             @NonNull PersistableBundle downloadParams) {
487                         Log.i(TAG, "Response : =" + downloadParams);
488                         statusLatch.countDown();
489                     }
490                 });
491         assertThat(statusLatch.await(2, SECONDS)).isTrue();
492 
493         // test download failed
494         Feature feature2 = new Feature.Builder(2).build();
495         CountDownLatch statusLatch2 = new CountDownLatch(1);
496 
497         mOnDeviceIntelligenceManager.requestFeatureDownload(feature2, null, EXECUTOR,
498                 new DownloadCallback() {
499                     @Override
500                     public void onDownloadFailed(int failureStatus,
501                             @Nullable String errorMessage,
502                             @NonNull PersistableBundle errorParams) {
503                         Log.e(TAG, "Got Error", new RuntimeException(errorMessage));
504                         statusLatch2.countDown();
505                     }
506 
507                     @Override
508                     public void onDownloadCompleted(
509                             @NonNull PersistableBundle downloadParams) {
510                         Log.i(TAG, "Response : =" + downloadParams);
511                     }
512                 });
513         assertThat(statusLatch2.await(2, SECONDS)).isTrue();
514     }
515 
516     @Test
517     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
resultPopulatedWhenRequestTokenInfo()518     public void resultPopulatedWhenRequestTokenInfo() throws Exception {
519         getInstrumentation()
520                 .getUiAutomation()
521                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
522         CountDownLatch statusLatch = new CountDownLatch(1);
523 
524         Feature feature = new Feature.Builder(1).build();
525         Bundle request = new Bundle();
526         request.putInt(TOKEN_INFO_COUNT_KEY, 0);
527         TokenInfo expectedTokenInfo = constructTokenInfo(0, null);
528         mOnDeviceIntelligenceManager.requestTokenInfo(feature, request
529                 , null,
530                 EXECUTOR,
531                 result -> {
532                     Log.i(TAG, "Response : =" + result.getCount());
533                     assertEquals(expectedTokenInfo.getCount(), result.getCount());
534                     statusLatch.countDown();
535                 });
536         assertThat(statusLatch.await(1, SECONDS)).isTrue();
537 
538 
539         PersistableBundle params = new PersistableBundle();
540         params.putInt("abc", 1);
541         request.putParcelable(TOKEN_INFO_PARAMS_KEY, params);
542         TokenInfo expectedTokenInfo2 = constructTokenInfo(0, params);
543         mOnDeviceIntelligenceManager.requestTokenInfo(feature, request
544                 , null,
545                 EXECUTOR,
546                 result -> {
547                     Log.i(TAG, "Response : =" + result.getCount());
548                     assertEquals(expectedTokenInfo2.getCount(), result.getCount());
549                     assertEquals(expectedTokenInfo2.getInfoParams().containsKey("abc"),
550                             result.getInfoParams().containsKey("abc"));
551                     statusLatch.countDown();
552                 });
553         assertThat(statusLatch.await(1, SECONDS)).isTrue();
554     }
555 
556     @Test
557     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
resultPopulatedWhenAttemptingProcessRequest()558     public void resultPopulatedWhenAttemptingProcessRequest() throws Exception {
559         getInstrumentation()
560                 .getUiAutomation()
561                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
562         CountDownLatch statusLatch = new CountDownLatch(1);
563         Feature feature = new Feature.Builder(1).build();
564         mOnDeviceIntelligenceManager.processRequest(feature,
565                 new Bundle(), 1, null,
566                 null, EXECUTOR, new ProcessingCallback() {
567                     @Override
568                     public void onResult(@NonNull Bundle result) {
569                         Log.i(TAG, "Final Result : " + result);
570                         statusLatch.countDown();
571                     }
572 
573                     @Override
574                     public void onError(@NonNull OnDeviceIntelligenceException error) {
575                         Log.e(TAG, "Error Occurred", error);
576                     }
577                 });
578         assertThat(statusLatch.await(1, SECONDS)).isTrue();
579     }
580 
581     @Test
582     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
resultPopulatedWhenAttemptingProcessRequestStreaming()583     public void resultPopulatedWhenAttemptingProcessRequestStreaming() throws Exception {
584         getInstrumentation()
585                 .getUiAutomation()
586                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
587         CountDownLatch statusLatch = new CountDownLatch(1);
588 
589         Feature feature = new Feature.Builder(1).build();
590         mOnDeviceIntelligenceManager.processRequestStreaming(feature,
591                 new Bundle(), 1,
592                 null, null, EXECUTOR,
593                 new StreamingProcessingCallback() {
594                     @Override
595                     public void onPartialResult(@NonNull Bundle partialResult) {
596                         Log.i(TAG, "New Content : " + partialResult);
597                     }
598 
599                     @Override
600                     public void onResult(Bundle result) {
601                         Log.i(TAG, "Final Result : " + result);
602                         statusLatch.countDown();
603                     }
604 
605                     @Override
606                     public void onError(@NonNull OnDeviceIntelligenceException error) {
607                         Log.e(TAG, "Final Result : ", error);
608                     }
609                 });
610         assertThat(statusLatch.await(1, SECONDS)).isTrue();
611     }
612 
613 
614 //===================== Tests Exception populated ==================
615 
616     @Test
617     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
exceptionPopulatedWhenAttemptingProcessRequest()618     public void exceptionPopulatedWhenAttemptingProcessRequest() throws Exception {
619         getInstrumentation()
620                 .getUiAutomation()
621                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
622         CountDownLatch statusLatch = new CountDownLatch(1);
623         Feature feature = new Feature.Builder(1).build();
624         Bundle bundle = new Bundle();
625         bundle.putInt(EXCEPTION_STATUS_CODE_KEY, 1);
626         OnDeviceIntelligenceException expectedException = constructException(bundle);
627         mOnDeviceIntelligenceManager.processRequest(feature, bundle, 1, null,
628                 null, EXECUTOR, new ProcessingCallback() {
629                     @Override
630                     public void onResult(@NonNull Bundle result) {
631                     }
632 
633                     @Override
634                     public void onError(@NonNull OnDeviceIntelligenceException error) {
635                         Log.e(TAG, "Error Occurred", error);
636                         assertEquals(error.getErrorCode(), expectedException.getErrorCode());
637                         statusLatch.countDown();
638                     }
639                 });
640         assertThat(statusLatch.await(1, SECONDS)).isTrue();
641 
642 
643         bundle.putString(EXCEPTION_MESSAGE_KEY, "test message");
644         OnDeviceIntelligenceException expectedException2 = constructException(bundle);
645         mOnDeviceIntelligenceManager.processRequest(feature, bundle, 1, null,
646                 null, EXECUTOR, new ProcessingCallback() {
647                     @Override
648                     public void onResult(@NonNull Bundle result) {
649                     }
650 
651                     @Override
652                     public void onError(@NonNull OnDeviceIntelligenceException error) {
653                         Log.e(TAG, "Error Occurred", error);
654                         assertEquals(error.getErrorCode(), expectedException2.getErrorCode());
655                         assertEquals(error.getMessage(), expectedException2.getMessage());
656                         statusLatch.countDown();
657                     }
658                 });
659         assertThat(statusLatch.await(1, SECONDS)).isTrue();
660 
661 
662         PersistableBundle params = new PersistableBundle();
663         params.putInt("abc", 1);
664         bundle.putParcelable(EXCEPTION_PARAMS_KEY, params);
665         OnDeviceIntelligenceException expectedException3 = constructException(bundle);
666         mOnDeviceIntelligenceManager.processRequest(feature, bundle, 1, null,
667                 null, EXECUTOR, new ProcessingCallback() {
668                     @Override
669                     public void onResult(@NonNull Bundle result) {
670                     }
671 
672                     @Override
673                     public void onError(@NonNull OnDeviceIntelligenceException error) {
674                         Log.e(TAG, "Error Occurred", error);
675                         assertEquals(error.getErrorCode(), expectedException3.getErrorCode());
676                         assertEquals(error.getMessage(), expectedException3.getMessage());
677                         assertEquals(error.getErrorParams().containsKey("abc"),
678                                 expectedException3.getErrorParams().containsKey("abc"));
679                         statusLatch.countDown();
680                     }
681                 });
682         assertThat(statusLatch.await(1, SECONDS)).isTrue();
683     }
684 
685 //===================== Tests for Processing and Cancellation signals  ==========================
686 
687     @Test
688     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
cancellationPropagatedWhenInvokedDuringRequest()689     public void cancellationPropagatedWhenInvokedDuringRequest() throws Exception {
690         getInstrumentation()
691                 .getUiAutomation()
692                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
693         CountDownLatch statusLatch = new CountDownLatch(2);
694         CancellationSignal cancellationSignal = new CancellationSignal();
695         Feature feature = new Feature.Builder(1).build();
696         CompletableFuture<Bundle> resultBundle = new CompletableFuture<>();
697         mOnDeviceIntelligenceManager.processRequestStreaming(feature,
698                 new Bundle(), 1, cancellationSignal,
699                 null, EXECUTOR, new StreamingProcessingCallback() {
700                     @Override
701                     public void onPartialResult(@NonNull Bundle partialResult) {
702                         Log.i(TAG, "New Content : " + partialResult);
703                         cancellationSignal.cancel(); //cancel
704                         statusLatch.countDown();
705                     }
706 
707                     @Override
708                     public void onResult(Bundle result) {
709                         Log.i(TAG, "Final Result : " + result);
710                         resultBundle.complete(result);
711                         statusLatch.countDown();
712                     }
713 
714                     @Override
715                     public void onError(@NonNull OnDeviceIntelligenceException error) {
716                         Log.e(TAG, "Final Result : ", error);
717                     }
718                 });
719         assertThat(statusLatch.await(2, SECONDS)).isTrue();
720         assertThat(resultBundle.get()).isNotNull();
721         assertThat(resultBundle.get().containsKey("test_key")).isTrue();
722         assertThat(resultBundle.get().getBoolean(TEST_KEY)).isTrue();
723     }
724 
725     @Test
726     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
cancellationPropagatedWhenInvokedBeforeMakingRequest()727     public void cancellationPropagatedWhenInvokedBeforeMakingRequest() throws Exception {
728         getInstrumentation()
729                 .getUiAutomation()
730                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
731         CountDownLatch statusLatch = new CountDownLatch(1);
732         CancellationSignal cancellationSignal = new CancellationSignal();
733         cancellationSignal.cancel(); //cancel
734         Feature feature = new Feature.Builder(1).build();
735         CompletableFuture<Bundle> resultBundle = new CompletableFuture<>();
736         mOnDeviceIntelligenceManager.processRequestStreaming(feature,
737                 new Bundle(), 1, cancellationSignal,
738                 null, EXECUTOR, new StreamingProcessingCallback() {
739                     @Override
740                     public void onPartialResult(@NonNull Bundle partialResult) {
741                         Log.i(TAG, "New Content : " + partialResult);
742                         statusLatch.countDown();
743                     }
744 
745                     @Override
746                     public void onResult(Bundle result) {
747                         Log.i(TAG, "Final Result : " + result);
748                         resultBundle.complete(result);
749                         statusLatch.countDown();
750                     }
751 
752                     @Override
753                     public void onError(@NonNull OnDeviceIntelligenceException error) {
754                         Log.e(TAG, "Final Result : ", error);
755                     }
756                 });
757         assertThat(statusLatch.await(1, SECONDS)).isTrue();
758         assertThat(resultBundle.get()).isNotNull();
759         assertThat(
760                 resultBundle.get().isEmpty()).isTrue(); // When cancelled before sending request,
761         // we simulate empty response.
762     }
763 
764     @Test
765     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
signalPropagatedWhenSignalIsInvokedBeforeAndDuringRequest()766     public void signalPropagatedWhenSignalIsInvokedBeforeAndDuringRequest() throws Exception {
767         getInstrumentation()
768                 .getUiAutomation()
769                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
770         CountDownLatch statusLatch = new CountDownLatch(4);
771         ProcessingSignal processingSignal = new ProcessingSignal();
772         processingSignal.sendSignal(PersistableBundle.EMPTY);
773         processingSignal.sendSignal(PersistableBundle.EMPTY);
774         Feature feature = new Feature.Builder(1).build();
775         mOnDeviceIntelligenceManager.processRequestStreaming(feature,
776                 new Bundle(), 1, null,
777                 processingSignal, EXECUTOR, new StreamingProcessingCallback() {
778                     @Override
779                     public void onPartialResult(@NonNull Bundle partialResult) {
780                         Log.i(TAG, "New Content : " + partialResult);
781                         statusLatch.countDown();
782                     }
783 
784                     @Override
785                     public void onResult(Bundle result) {
786                         Log.i(TAG, "Final Result : " + result);
787                         statusLatch.countDown();
788                     }
789 
790                     @Override
791                     public void onError(@NonNull OnDeviceIntelligenceException error) {
792                         Log.e(TAG, "Final Result : ", error);
793                     }
794                 });
795         processingSignal.sendSignal(PersistableBundle.EMPTY);
796         assertThat(statusLatch.await(2, SECONDS)).isTrue();
797     }
798 
799     //===================== Tests for Manager Methods When No Service is Configured =============
800 
801     @Test
802     @SkipSetupAndTeardown
803     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
exceptionWhenAttemptingGetVersionWithoutServiceConfigured()804     public void exceptionWhenAttemptingGetVersionWithoutServiceConfigured() {
805         getInstrumentation()
806                 .getUiAutomation()
807                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
808         assumeFalse("Service is already configured as part of the device overlay config.",
809                 isServiceOverlayConfigured());
810         mOnDeviceIntelligenceManager =
811                 (OnDeviceIntelligenceManager)
812                         mContext.getSystemService(Context.ON_DEVICE_INTELLIGENCE_SERVICE);
813         clearTestableOnDeviceIntelligenceService();
814         // Test throws IllegalStateException
815         assertThrows("no service configured to perform getVersion",
816                 IllegalStateException.class,
817                 () -> mOnDeviceIntelligenceManager.getVersion(EXECUTOR,
818                         result -> Log.i(TAG, "Feature : =" + result)));
819     }
820 
821     @Test
822     @SkipSetupAndTeardown
823     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
exceptionWhenAttemptingProcessRequestWithoutServiceConfigured()824     public void exceptionWhenAttemptingProcessRequestWithoutServiceConfigured() {
825         getInstrumentation()
826                 .getUiAutomation()
827                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
828         assumeFalse("Service is already configured as part of the device overlay config.",
829                 isServiceOverlayConfigured());
830         mOnDeviceIntelligenceManager =
831                 (OnDeviceIntelligenceManager)
832                         mContext.getSystemService(Context.ON_DEVICE_INTELLIGENCE_SERVICE);
833         clearTestableOnDeviceIntelligenceService();
834         Feature feature = new Feature.Builder(1).build();
835         // Test throws IllegalStateException
836         assertThrows(
837                 "no service configured for processRequestStreaming",
838                 IllegalStateException.class,
839                 () -> mOnDeviceIntelligenceManager.processRequestStreaming(feature,
840                         new Bundle(), 1,
841                         null, null, EXECUTOR,
842                         new StreamingProcessingCallback() {
843                             @Override
844                             public void onPartialResult(@NonNull Bundle partialResult) {
845                                 Log.i(TAG, "New Content : " + partialResult);
846                             }
847 
848                             @Override
849                             public void onResult(Bundle result) {
850                                 Log.i(TAG, "Final Result : " + result);
851                             }
852 
853                             @Override
854                             public void onError(@NonNull OnDeviceIntelligenceException error) {
855                                 Log.e(TAG, "Final Result : ", error);
856                             }
857                         }));
858     }
859 
860     // ========= Test package manager returns parent process package name for isolated_compute_app
861     @Test
862     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
inferenceServiceShouldReturnParentPackageName()863     public void inferenceServiceShouldReturnParentPackageName() throws Exception {
864         getInstrumentation()
865                 .getUiAutomation()
866                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
867         CountDownLatch statusLatch = new CountDownLatch(1);
868         Feature feature = new Feature.Builder(1).build();
869         CompletableFuture<String> packageNameFuture = new CompletableFuture<>();
870         mOnDeviceIntelligenceManager.processRequest(feature,
871                 Bundle.EMPTY, REQUEST_TYPE_GET_PACKAGE_NAME, null,
872                 null, EXECUTOR, new ProcessingCallback() {
873                     @Override
874                     public void onResult(@NonNull Bundle result) {
875                         Log.i(TAG, "Final Result : " + result);
876                         packageNameFuture.complete(result.getString(TEST_KEY));
877                         statusLatch.countDown();
878                     }
879 
880                     @Override
881                     public void onError(@NonNull OnDeviceIntelligenceException error) {
882                         Log.e(TAG, "Error Occurred", error);
883                     }
884                 });
885         assertThat(statusLatch.await(1, SECONDS)).isTrue();
886         assertThat(packageNameFuture.get()).isEqualTo(CTS_PACKAGE_NAME);
887     }
888 
889     @Test
890     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
callerUidReceivedIsOriginalCallerUid()891     public void callerUidReceivedIsOriginalCallerUid() throws Exception {
892         getInstrumentation()
893                 .getUiAutomation()
894                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
895         CountDownLatch statusLatch = new CountDownLatch(1);
896         Feature feature = new Feature.Builder(1).build();
897         mOnDeviceIntelligenceManager.processRequest(feature,
898                 Bundle.EMPTY, REQUEST_TYPE_GET_CALLER_UID, null,
899                 null, EXECUTOR, new ProcessingCallback() {
900                     @Override
901                     public void onResult(@NonNull Bundle result) {
902                         Log.i(TAG, "Final Result : " + result);
903                         assertThat(result.getInt(TEST_KEY)).isEqualTo(Process.myUid());
904                         statusLatch.countDown();
905                     }
906 
907                     @Override
908                     public void onError(@NonNull OnDeviceIntelligenceException error) {
909                         Log.e(TAG, "Error Occurred", error);
910                     }
911                 });
912         assertThat(statusLatch.await(1, SECONDS)).isTrue();
913     }
914 
915 
916     //===================== Tests for accessing file from isolated process via non-isolated =======
917     @Test
918     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
canAccessFilesInIsolated()919     public void canAccessFilesInIsolated() throws Exception {
920         int[] requestTypes =
921                 new int[]{REQUEST_TYPE_GET_FILE_FROM_MAP, REQUEST_TYPE_GET_FILE_FROM_STREAM,
922                         REQUEST_TYPE_GET_FILE_FROM_PFD,
923                         REQUEST_TYPE_GET_FILE_FROM_NON_FILES_DIRECTORY};
924         for (int requestType : requestTypes) {
925             sendRequestToReadTestFile(requestType);
926         }
927     }
928 
sendRequestToReadTestFile(int requestType)929     private void sendRequestToReadTestFile(int requestType)
930             throws InterruptedException, ExecutionException {
931         getInstrumentation()
932                 .getUiAutomation()
933                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
934         Feature feature = new Feature.Builder(1).build();
935         CountDownLatch statusLatch = new CountDownLatch(1);
936         CompletableFuture<String> fileContents = new CompletableFuture<>();
937         mOnDeviceIntelligenceManager.processRequest(feature,
938                 Bundle.EMPTY, requestType, null,
939                 null, EXECUTOR, new StreamingProcessingCallback() {
940                     @Override
941                     public void onPartialResult(@NonNull Bundle partialResult) {
942                         Log.i(TAG, "New Content : " + partialResult);
943                     }
944 
945                     @Override
946                     public void onResult(Bundle result) {
947                         Log.i(TAG, "Final Result : " + result);
948                         fileContents.complete(result.getString(TEST_KEY));
949                         statusLatch.countDown();
950                     }
951 
952                     @Override
953                     public void onError(@NonNull OnDeviceIntelligenceException error) {
954                         Log.e(TAG, "Final Result : ", error);
955                     }
956                 });
957         assertThat(statusLatch.await(1, SECONDS)).isTrue();
958         assertThat(fileContents.get()).isEqualTo(TEST_CONTENT);
959     }
960 
961     @Test
962     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
updateProcessingStateReturnsSuccessfully()963     public void updateProcessingStateReturnsSuccessfully() throws Exception {
964         // When targets run as a different user than 0, it is not possible to get service
965         // instance from user 0 in this test.
966         assumeTrue(isSystemUser());
967         getInstrumentation()
968                 .getUiAutomation()
969                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
970         CountDownLatch statusLatch = new CountDownLatch(1);
971         // init the intelligence service
972         CtsIntelligenceService.initServiceConnectionLatch();
973         mOnDeviceIntelligenceManager.getVersion(EXECUTOR, unused -> statusLatch.countDown());
974         statusLatch.await(1, SECONDS);
975 
976         // call update state on the service instance
977         CtsIntelligenceService.waitServiceConnect();
978         OnDeviceIntelligenceService onDeviceIntelligenceService =
979                 CtsIntelligenceService.getServiceInstance();
980         CountDownLatch statusLatch2 = new CountDownLatch(1);
981         onDeviceIntelligenceService.updateProcessingState(Bundle.EMPTY, EXECUTOR, result -> {
982             assertThat(result.isEmpty()).isTrue();
983             statusLatch2.countDown();
984         });
985 
986         assertThat(statusLatch2.await(1, SECONDS)).isTrue();
987     }
988 
989 
990     @Test
991     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE_MODULE)
getLatestInferenceInfoReturnSuccessfully()992     public void getLatestInferenceInfoReturnSuccessfully() throws Exception {
993         // When targets run as a different user than 0, it is not possible to get service
994         // instance from user 0 in this test.
995         assumeTrue(isSystemUser());
996         getInstrumentation()
997                 .getUiAutomation()
998                 .adoptShellPermissionIdentity(Manifest.permission.DUMP,
999                         Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
1000         CountDownLatch statusLatch = new CountDownLatch(1);
1001         mOnDeviceIntelligenceManager.processRequest(new Feature.Builder(1).build(),
1002                 Bundle.EMPTY, REQUEST_TYPE_POPULATE_INFERENCE_INFO, null,
1003                 null, EXECUTOR, new ProcessingCallback() {
1004                     @Override
1005                     public void onResult(@NonNull Bundle result) {
1006                         Log.i(TAG, "Final Result : " + result);
1007                         statusLatch.countDown();
1008                     }
1009 
1010                     @Override
1011                     public void onError(@NonNull OnDeviceIntelligenceException error) {
1012                         Log.e(TAG, "Error Occurred", error);
1013                     }
1014                 });
1015         assertThat(statusLatch.await(1, SECONDS)).isTrue();
1016         List<InferenceInfo> inferenceInfoList = mOnDeviceIntelligenceManager.getLatestInferenceInfo(
1017                 0);
1018         assertThat(inferenceInfoList).isNotEmpty();
1019     }
1020 
1021     //===================== Tests data augmentation while processing request =====================
1022     @Test
1023     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
dataAugmentationReturnsDataToInference()1024     public void dataAugmentationReturnsDataToInference() throws Exception {
1025         getInstrumentation()
1026                 .getUiAutomation()
1027                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
1028         Feature feature = new Feature.Builder(1).build();
1029         CountDownLatch statusLatch = new CountDownLatch(1);
1030         CompletableFuture<String> augmentedContent = new CompletableFuture<>();
1031         mOnDeviceIntelligenceManager.processRequest(feature,
1032                 Bundle.EMPTY, REQUEST_TYPE_GET_AUGMENTED_DATA, null,
1033                 null, EXECUTOR, new StreamingProcessingCallback() {
1034                     @Override
1035                     public void onPartialResult(@NonNull Bundle partialResult) {
1036                         Log.i(TAG, "New Content : " + partialResult);
1037                     }
1038 
1039                     @Override
1040                     public void onResult(Bundle result) {
1041                         Log.i(TAG, "Final Result : " + result);
1042                         augmentedContent.complete(result.getString(TEST_AUGMENT_KEY));
1043                         statusLatch.countDown();
1044                     }
1045 
1046                     @Override
1047                     public void onError(@NonNull OnDeviceIntelligenceException error) {
1048                         Log.e(TAG, "Final Result : ", error);
1049                     }
1050 
1051                     @Override
1052                     public void onDataAugmentRequest(Bundle processedContent,
1053                             Consumer<Bundle> contentConsumer) {
1054                         Bundle bundle = new Bundle();
1055                         bundle.putString(TEST_AUGMENT_KEY, TEST_AUGMENT_CONTENT);
1056                         contentConsumer.accept(bundle);
1057                     }
1058                 });
1059         assertThat(statusLatch.await(1, SECONDS)).isTrue();
1060         assertThat(augmentedContent.get()).isEqualTo(TEST_AUGMENT_CONTENT);
1061     }
1062 
1063     //===================== Tests broadcasts are sent for model updates =========================
1064     @Test
1065     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
broadcastsMustBeSentOnModelUpdates()1066     public void broadcastsMustBeSentOnModelUpdates() throws Exception {
1067         assumeTrue(isSystemUser());
1068         getInstrumentation()
1069                 .getUiAutomation()
1070                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE);
1071         setTestableBroadcastKeys(new String[]{MODEL_LOADED_BROADCAST_ACTION, "blah"},
1072                 mContext.getPackageName());
1073         Feature feature = new Feature.Builder(1).build();
1074         CountDownLatch statusLatch = new CountDownLatch(2);
1075         BroadcastReceiver broadcastReceiver = new BroadcastReceiver() {
1076             @Override
1077             public void onReceive(Context context, Intent intent) {
1078                 String action = intent.getAction();
1079                 if (action != null) {
1080                     Log.d(TAG, "Received broadcast with action: " + action);
1081                     if (action == MODEL_LOADED_BROADCAST_ACTION) {
1082                         statusLatch.countDown();
1083                     }
1084                 }
1085             }
1086         };
1087         mContext.registerReceiver(broadcastReceiver,
1088                 new IntentFilter(MODEL_LOADED_BROADCAST_ACTION), RECEIVER_EXPORTED);
1089         mOnDeviceIntelligenceManager.processRequest(feature,
1090                 Bundle.EMPTY, 1, null,
1091                 null, EXECUTOR, new StreamingProcessingCallback() {
1092                     @Override
1093                     public void onPartialResult(@NonNull Bundle partialResult) {
1094                         Log.i(TAG, "New Content : " + partialResult);
1095                     }
1096 
1097                     @Override
1098                     public void onResult(Bundle result) {
1099                         Log.i(TAG, "Final Result : " + result);
1100                         statusLatch.countDown();
1101                     }
1102 
1103                     @Override
1104                     public void onError(@NonNull OnDeviceIntelligenceException error) {
1105                         Log.e(TAG, "Final Result : ", error);
1106                     }
1107                 });
1108         assertThat(statusLatch.await(5, SECONDS)).isTrue();
1109     }
1110 
1111     //===================== Tests unbind based on timeout settings are invoked ====================
1112 
1113     @Test
1114     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
serviceUnbindsWhenCallbackIsNotPopulatedAfterIdleTimeout()1115     public void serviceUnbindsWhenCallbackIsNotPopulatedAfterIdleTimeout() throws Exception {
1116         getInstrumentation()
1117                 .getUiAutomation()
1118                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE,
1119                         Manifest.permission.WRITE_SECURE_SETTINGS);
1120         assumeTrue(isSystemUser());
1121         updateSecureSettings();
1122         // Feature Id to ensure no callbacks are invoked
1123         Feature feature = new Feature.Builder(3).build();
1124         CtsIntelligenceService.initServiceConnectionLatch();
1125         CtsIntelligenceService.initUnbindLatch();
1126         mOnDeviceIntelligenceManager.requestFeatureDownload(feature, null, EXECUTOR,
1127                 new DownloadCallback() {
1128                     @Override
1129                     public void onDownloadFailed(int failureStatus,
1130                             @Nullable String errorMessage,
1131                             @NonNull PersistableBundle errorParams) {
1132                         Log.e(TAG, "Got Error", new RuntimeException(errorMessage));
1133                     }
1134 
1135                     @Override
1136                     public void onDownloadProgress(long bytesDownloaded) {
1137                     }
1138 
1139                     @Override
1140                     public void onDownloadStarted(long bytesDownloaded) {
1141 
1142                     }
1143 
1144                     @Override
1145                     public void onDownloadCompleted(
1146                             @NonNull PersistableBundle downloadParams) {
1147                         Log.i(TAG, "Response : =" + downloadParams);
1148                     }
1149                 });
1150         CtsIntelligenceService.waitServiceConnect();
1151         CtsIntelligenceService.waitForUnbind();
1152         resetSecureSettings();
1153     }
1154 
1155     @Test
1156     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
serviceUnbindsWhenCallbackIsPopulatedAfterIdleTimeout()1157     public void serviceUnbindsWhenCallbackIsPopulatedAfterIdleTimeout() throws Exception {
1158         getInstrumentation()
1159                 .getUiAutomation()
1160                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE,
1161                         Manifest.permission.WRITE_SECURE_SETTINGS);
1162         assumeTrue(isSystemUser());
1163         updateSecureSettings();
1164         CtsIntelligenceService.initServiceConnectionLatch();
1165         CtsIntelligenceService.initUnbindLatch();
1166         CountDownLatch statusLatch = new CountDownLatch(1);
1167 
1168         mOnDeviceIntelligenceManager.getVersion(EXECUTOR,
1169                 result -> {
1170                     Log.i(TAG, "Version : =" + result);
1171                     statusLatch.countDown();
1172                 });
1173         assertThat(statusLatch.await(1, SECONDS)).isTrue();
1174         CtsIntelligenceService.waitServiceConnect();
1175         CtsIntelligenceService.waitForUnbind();
1176         resetSecureSettings();
1177     }
1178 
1179     @Test
1180     @RequiresFlagsEnabled(FLAG_ENABLE_ON_DEVICE_INTELLIGENCE)
deviceConfigUpdateMustBeSentOnInferenceServiceConnected()1181     public void deviceConfigUpdateMustBeSentOnInferenceServiceConnected() throws Exception {
1182         getInstrumentation()
1183                 .getUiAutomation()
1184                 .adoptShellPermissionIdentity(Manifest.permission.USE_ON_DEVICE_INTELLIGENCE,
1185                         "android.permission.WRITE_DEVICE_CONFIG",
1186                         "android.permission.WRITE_ALLOWLISTED_DEVICE_CONFIG",
1187                         "android.permission.READ_DEVICE_CONFIG",
1188                         "android.permission.MONITOR_DEVICE_CONFIG_ACCESS");
1189         Feature feature = new Feature.Builder(1).build();
1190         CountDownLatch statusLatch = new CountDownLatch(1);
1191         String currentVal = DeviceConfig.getProperty(TEST_OD_NAMESPACE, "key1");
1192         if (currentVal == null) {
1193             currentVal = "val1";
1194         }
1195         String modifiedVal = currentVal + "_new";
1196         mOnDeviceIntelligenceManager.processRequest(feature,
1197                 Bundle.EMPTY, REQUEST_TYPE_GET_UPDATED_DEVICE_CONFIG, null,
1198                 null, EXECUTOR, new ProcessingCallback() {
1199                     @Override
1200                     public void onResult(Bundle result) {
1201                         Log.i(TAG, "Final Result : " + result);
1202                         PersistableBundle receivedConfig = result.getParcelable(TEST_KEY,
1203                                 PersistableBundle.class);
1204                         assertThat(receivedConfig.containsKey("key1")).isTrue();
1205                         assertThat(receivedConfig.getString("key1")).isEqualTo(modifiedVal);
1206 
1207                         statusLatch.countDown();
1208                     }
1209 
1210                     @Override
1211                     public void onError(@NonNull OnDeviceIntelligenceException error) {
1212                         Log.e(TAG, "Final Result : ", error);
1213                     }
1214                 });
1215         Executors.newScheduledThreadPool(1).schedule(
1216                 () -> {
1217                     DeviceConfig.setProperty(TEST_OD_NAMESPACE, "key1", modifiedVal, false);
1218                     Log.i(TAG, "Finished writing property to device config.");
1219                 }, 2L,
1220                 SECONDS);
1221         assertThat(statusLatch.await(10, SECONDS)).isTrue();
1222         DeviceConfig.deleteProperty(TEST_OD_NAMESPACE, "key1");
1223     }
1224 
1225 
clearTestableOnDeviceIntelligenceService()1226     public static void clearTestableOnDeviceIntelligenceService() {
1227         runShellCommand("cmd on_device_intelligence set-temporary-services");
1228     }
1229 
bindToTestableOnDeviceIntelligenceServices()1230     public void bindToTestableOnDeviceIntelligenceServices() {
1231         setTestableOnDeviceIntelligenceServiceNames(
1232                 new String[]{CTS_INTELLIGENCE_SERVICE_NAME, CTS_INFERENCE_SERVICE_NAME});
1233         assertThat(CTS_INFERENCE_SERVICE_NAME).contains(getOnDeviceIntelligencePackageName());
1234     }
1235 
updateSecureSettings()1236     private void updateSecureSettings() {
1237         Settings.Secure.putLong(mContext.getContentResolver(),
1238                 Settings.Secure.ON_DEVICE_INTELLIGENCE_UNBIND_TIMEOUT_MS, SECONDS.toMillis(1));
1239         Settings.Secure.putLong(mContext.getContentResolver(),
1240                 Settings.Secure.ON_DEVICE_INTELLIGENCE_IDLE_TIMEOUT_MS, SECONDS.toMillis(1));
1241     }
1242 
resetSecureSettings()1243     private void resetSecureSettings() {
1244         Settings.Secure.putLong(mContext.getContentResolver(),
1245                 Settings.Secure.ON_DEVICE_INTELLIGENCE_UNBIND_TIMEOUT_MS, -1);
1246         Settings.Secure.putLong(mContext.getContentResolver(),
1247                 Settings.Secure.ON_DEVICE_INTELLIGENCE_IDLE_TIMEOUT_MS, HOURS.toMillis(1));
1248     }
1249 
getOnDeviceIntelligencePackageName()1250     private String getOnDeviceIntelligencePackageName() {
1251         return mOnDeviceIntelligenceManager.getRemoteServicePackageName();
1252     }
1253 
isServiceOverlayConfigured()1254     private boolean isServiceOverlayConfigured() {
1255         String sanboxedServiceComponentName = mContext.getResources()
1256                         .getString(
1257                                 mContext.getResources()
1258                                         .getIdentifier(
1259                                                 "config_defaultOnDeviceSandboxedInferenceService",
1260                                                 "string",
1261                                                 "android"));
1262         String intelligenceServiceComponentName = mContext.getResources()
1263                         .getString(
1264                                 mContext.getResources()
1265                                         .getIdentifier(
1266                                                 "config_defaultOnDeviceIntelligenceService",
1267                                                 "string",
1268                                                 "android"));
1269 
1270         return !TextUtils.isEmpty(sanboxedServiceComponentName) || !TextUtils.isEmpty(
1271                 intelligenceServiceComponentName);
1272     }
1273 
isSystemUser()1274     private static boolean isSystemUser() {
1275         return Process.myUserHandle().equals(UserHandle.SYSTEM);
1276     }
1277 
setTestableBroadcastKeys(String[] broadcastKeys, String packageName)1278     public static void setTestableBroadcastKeys(String[] broadcastKeys, String packageName) {
1279         runShellCommand(
1280                 "cmd on_device_intelligence set-model-broadcasts %s %s %s %d",
1281                 broadcastKeys[0], broadcastKeys[1], packageName, TEMPORARY_SERVICE_DURATION);
1282     }
1283 
1284 
setTestableDeviceConfigNamespace(String configNamespace)1285     public static void setTestableDeviceConfigNamespace(String configNamespace) {
1286         runShellCommand(
1287                 "cmd on_device_intelligence set-deviceconfig-namespace %s %d", configNamespace,
1288                 TEMPORARY_SERVICE_DURATION);
1289     }
1290 
setTestableOnDeviceIntelligenceServiceNames(String[] serviceNames)1291     public static void setTestableOnDeviceIntelligenceServiceNames(String[] serviceNames) {
1292         runShellCommand(
1293                 "cmd on_device_intelligence set-temporary-services %s %s %d",
1294                 serviceNames[0], serviceNames[1], TEMPORARY_SERVICE_DURATION);
1295     }
1296 
1297     @Retention(RetentionPolicy.RUNTIME)
1298     @Target(ElementType.METHOD)
1299     public @interface SkipSetupAndTeardown {
1300     }
1301 
1302     @Rule
1303     public TestRule skipSetupAndTeardownRule = (base, description) -> new Statement() {
1304         @Override
1305         public void evaluate() throws Throwable {
1306             if (description.getAnnotation(SkipSetupAndTeardown.class) != null) {
1307                 // Skip setup and teardown for annotated tests
1308                 base.evaluate();
1309             } else {
1310                 // Run setup and teardown for other tests
1311                 setUp();
1312                 try {
1313                     base.evaluate();
1314                 } finally {
1315                     tearDown();
1316                 }
1317             }
1318         }
1319     };
1320 
1321 
1322 }
1323