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 com.android.server.ondeviceintelligence;
18 
19 import static com.google.common.truth.Truth.assertThat;
20 
21 import android.os.Bundle;
22 import android.os.PersistableBundle;
23 import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
24 import android.app.ondeviceintelligence.InferenceInfo;
25 import android.util.Base64;
26 
27 import androidx.test.ext.junit.runners.AndroidJUnit4;
28 
29 import com.android.framework.protobuf.nano.MessageNano;
30 
31 import org.junit.Before;
32 import org.junit.Test;
33 import org.junit.runner.RunWith;
34 
35 import java.io.IOException;
36 import java.util.List;
37 
38 @RunWith(AndroidJUnit4.class)
39 public class InferenceInfoStoreTest {
40     InferenceInfoStore inferenceInfoStore;
41 
42     @Before
setUp()43     public void setUp() {
44         inferenceInfoStore = new InferenceInfoStore(1000);
45     }
46 
47     @Test
testInferenceInfoParsesFromBundleSuccessfully()48     public void testInferenceInfoParsesFromBundleSuccessfully() throws Exception {
49         Bundle bundle = new Bundle();
50         bundle.putByteArray(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY,
51                 getInferenceInfoBytes(1, 1, 100));
52         inferenceInfoStore.addInferenceInfoFromBundle(bundle);
53         List<InferenceInfo> inferenceInfos = inferenceInfoStore.getLatestInferenceInfo(0);
54         assertThat(inferenceInfos).hasSize(1);
55         assertThat(inferenceInfos.get(0).getUid()).isEqualTo(1);
56         assertThat(inferenceInfos.get(0).getStartTimeMs()).isEqualTo(1);
57         assertThat(inferenceInfos.get(0).getEndTimeMs()).isEqualTo(100);
58     }
59 
60     @Test
testInferenceInfoParsesFromPersistableBundleSuccessfully()61     public void testInferenceInfoParsesFromPersistableBundleSuccessfully() throws Exception {
62         PersistableBundle bundle = new PersistableBundle();
63         bundle.putString(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY,
64                 Base64.encodeToString(getInferenceInfoBytes(1, 1, 100), Base64.DEFAULT));
65         inferenceInfoStore.addInferenceInfoFromBundle(bundle);
66         List<InferenceInfo> inferenceInfos = inferenceInfoStore.getLatestInferenceInfo(0);
67         assertThat(inferenceInfos).hasSize(1);
68         assertThat(inferenceInfos.get(0).getUid()).isEqualTo(1);
69         assertThat(inferenceInfos.get(0).getStartTimeMs()).isEqualTo(1);
70         assertThat(inferenceInfos.get(0).getEndTimeMs()).isEqualTo(100);
71     }
72 
73 
74     @Test
testEvictionAfterMaxAge()75     public void testEvictionAfterMaxAge() throws Exception {
76         PersistableBundle bundle = new PersistableBundle();
77         long testStartTime = System.currentTimeMillis();
78         bundle.putString(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY,
79                 Base64.encodeToString(getInferenceInfoBytes(1,  testStartTime - 10,
80                         testStartTime + 100), Base64.DEFAULT));
81         inferenceInfoStore.addInferenceInfoFromBundle(bundle);
82         bundle.putString(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY,
83                 Base64.encodeToString(getInferenceInfoBytes(1, testStartTime - 5,
84                         testStartTime + 100), Base64.DEFAULT));
85         inferenceInfoStore.addInferenceInfoFromBundle(bundle);
86         Thread.sleep(1020);
87         List<InferenceInfo> inferenceInfos = inferenceInfoStore.getLatestInferenceInfo(0);
88         assertThat(inferenceInfos).hasSize(2);
89         assertThat(inferenceInfos.get(0).getUid()).isEqualTo(1);
90         assertThat(inferenceInfos.get(0).getStartTimeMs()).isEqualTo(testStartTime - 10);
91         assertThat(inferenceInfos.get(0).getEndTimeMs()).isEqualTo(testStartTime + 100);
92         inferenceInfoStore.addInferenceInfoFromBundle(bundle);
93         List<InferenceInfo> inferenceInfos2 = inferenceInfoStore.getLatestInferenceInfo(0);
94         assertThat(inferenceInfos2).hasSize(1); //previous entries should have been evicted
95     }
96 
getInferenceInfoBytes(int uid, long startTime, long endTime)97     private byte[] getInferenceInfoBytes(int uid, long startTime, long endTime) {
98         com.android.server.ondeviceintelligence.nano.InferenceInfo inferenceInfo =
99                 new com.android.server.ondeviceintelligence.nano.InferenceInfo();
100         inferenceInfo.uid = uid;
101         inferenceInfo.startTimeMs = startTime;
102         inferenceInfo.endTimeMs = endTime;
103         return MessageNano.toByteArray(inferenceInfo);
104     }
105 }
106