Adding tests for InferenceInfoStore

Bug: 335390745
Change-Id: I913a9708d9e5b33f9b8a107c0bda1597b3ff047e
diff --git a/services/core/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java b/services/core/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
index 6578853..b532d5a 100644
--- a/services/core/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
+++ b/services/core/java/com/android/server/ondeviceintelligence/InferenceInfoStore.java
@@ -21,9 +21,9 @@
 import android.os.PersistableBundle;
 import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
 import android.util.Slog;
+import android.util.Base64;
 
 import java.io.IOException;
-import java.util.Base64;
 import java.util.Comparator;
 import java.util.List;
 import java.util.TreeSet;
@@ -53,7 +53,7 @@
             String infoBytesBase64String = pb.getString(
                     OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY);
             if (infoBytesBase64String != null) {
-                byte[] infoBytes = Base64.getDecoder().decode(infoBytesBase64String);
+                byte[] infoBytes = Base64.decode(infoBytesBase64String, Base64.DEFAULT);
                 com.android.server.ondeviceintelligence.nano.InferenceInfo inferenceInfo =
                         com.android.server.ondeviceintelligence.nano.InferenceInfo.parseFrom(
                                 infoBytes);
@@ -84,7 +84,9 @@
     }
 
     private synchronized void add(com.android.server.ondeviceintelligence.nano.InferenceInfo info) {
-        while (System.currentTimeMillis() - inferenceInfos.first().getStartTimeMs() > maxAgeMs) {
+        while (!inferenceInfos.isEmpty()
+                && System.currentTimeMillis() - inferenceInfos.first().getStartTimeMs()
+                > maxAgeMs) {
             inferenceInfos.pollFirst();
         }
         inferenceInfos.add(toInferenceInfo(info));
diff --git a/services/tests/ondeviceintelligencetests/Android.bp b/services/tests/ondeviceintelligencetests/Android.bp
new file mode 100644
index 0000000..aa85942
--- /dev/null
+++ b/services/tests/ondeviceintelligencetests/Android.bp
@@ -0,0 +1,62 @@
+// Copyright (C) 2024 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package {
+    // See: http://go/android-license-faq
+    // A large-scale-change added 'default_applicable_licenses' to import
+    // all of the 'license_kinds' from "frameworks_base_license"
+    // to get the below license kinds:
+    //   SPDX-license-identifier-Apache-2.0
+    default_applicable_licenses: ["frameworks_base_license"],
+}
+
+android_test {
+    name: "FrameworksOnDeviceIntelligenceTests",
+    team: "trendy_team_machine_learning",
+    defaults: [
+        "modules-utils-testable-device-config-defaults",
+    ],
+
+    srcs: [
+        "src/**/*.java",
+    ],
+
+    static_libs: [
+        "androidx.test.core",
+        "androidx.test.runner",
+        "androidx.test.ext.truth",
+        "mockito-target-extended-minus-junit4",
+        "platform-test-annotations",
+        "services.core",
+        "servicestests-core-utils",
+        "servicestests-utils-mockito-extended",
+        "truth",
+        "frameworks-base-testutils",
+        "androidx.test.rules",
+    ],
+
+    libs: [
+        "android.test.mock",
+        "android.test.base",
+        "android.test.runner",
+    ],
+
+    certificate: "platform",
+    platform_apis: true,
+    test_suites: ["device-tests"],
+
+    optimize: {
+        enabled: false,
+    },
+}
diff --git a/services/tests/ondeviceintelligencetests/AndroidManifest.xml b/services/tests/ondeviceintelligencetests/AndroidManifest.xml
new file mode 100644
index 0000000..8bd111e
--- /dev/null
+++ b/services/tests/ondeviceintelligencetests/AndroidManifest.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+          package="com.android.frameworks.ondeviceintelligencetests">
+
+    <application android:testOnly="true"
+                 android:debuggable="true">
+        <uses-library android:name="android.test.runner" />
+    </application>
+
+    <instrumentation
+        android:name="androidx.test.runner.AndroidJUnitRunner"
+        android:targetPackage="com.android.frameworks.ondeviceintelligencetests"
+        android:label="Frameworks OnDeviceIntelligence Services Tests" />
+
+</manifest>
diff --git a/services/tests/ondeviceintelligencetests/AndroidTest.xml b/services/tests/ondeviceintelligencetests/AndroidTest.xml
new file mode 100644
index 0000000..3ae96c5
--- /dev/null
+++ b/services/tests/ondeviceintelligencetests/AndroidTest.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2024 The Android Open Source Project
+
+     Licensed under the Apache License, Version 2.0 (the "License");
+     you may not use this file except in compliance with the License.
+     You may obtain a copy of the License at
+
+          http://www.apache.org/licenses/LICENSE-2.0
+
+     Unless required by applicable law or agreed to in writing, software
+     distributed under the License is distributed on an "AS IS" BASIS,
+     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+     See the License for the specific language governing permissions and
+     limitations under the License.
+-->
+<configuration description="Runs Frameworks OnDeviceIntelligence Services Tests.">
+    <target_preparer class="com.android.tradefed.targetprep.suite.SuiteApkInstaller">
+        <option name="cleanup-apks" value="true" />
+        <option name="install-arg" value="-t" />
+        <option name="test-file-name" value="FrameworksOnDeviceIntelligenceTests.apk" />
+    </target_preparer>
+
+    <option name="test-tag" value="FrameworksOnDeviceIntelligenceTests" />
+
+    <test class="com.android.tradefed.testtype.AndroidJUnitTest" >
+        <option name="package" value="com.android.frameworks.ondeviceintelligencetests" />
+        <option name="runner" value="androidx.test.runner.AndroidJUnitRunner" />
+        <option name="hidden-api-checks" value="false"/>
+    </test>
+</configuration>
diff --git a/services/tests/ondeviceintelligencetests/src/com/android/server/ondeviceintelligence/InferenceInfoStoreTest.java b/services/tests/ondeviceintelligencetests/src/com/android/server/ondeviceintelligence/InferenceInfoStoreTest.java
new file mode 100644
index 0000000..d3943e3
--- /dev/null
+++ b/services/tests/ondeviceintelligencetests/src/com/android/server/ondeviceintelligence/InferenceInfoStoreTest.java
@@ -0,0 +1,105 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.ondeviceintelligence;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import android.os.Bundle;
+import android.os.PersistableBundle;
+import android.service.ondeviceintelligence.OnDeviceSandboxedInferenceService;
+import android.app.ondeviceintelligence.InferenceInfo;
+import android.util.Base64;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+
+import com.android.framework.protobuf.nano.MessageNano;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import java.io.IOException;
+import java.util.List;
+
+@RunWith(AndroidJUnit4.class)
+public class InferenceInfoStoreTest {
+    InferenceInfoStore inferenceInfoStore;
+
+    @Before
+    public void setUp() {
+        inferenceInfoStore = new InferenceInfoStore(1000);
+    }
+
+    @Test
+    public void testInferenceInfoParsesFromBundleSuccessfully() throws Exception {
+        Bundle bundle = new Bundle();
+        bundle.putByteArray(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY,
+                getInferenceInfoBytes(1, 1, 100));
+        inferenceInfoStore.addInferenceInfoFromBundle(bundle);
+        List<InferenceInfo> inferenceInfos = inferenceInfoStore.getLatestInferenceInfo(0);
+        assertThat(inferenceInfos).hasSize(1);
+        assertThat(inferenceInfos.get(0).getUid()).isEqualTo(1);
+        assertThat(inferenceInfos.get(0).getStartTimeMs()).isEqualTo(1);
+        assertThat(inferenceInfos.get(0).getEndTimeMs()).isEqualTo(100);
+    }
+
+    @Test
+    public void testInferenceInfoParsesFromPersistableBundleSuccessfully() throws Exception {
+        PersistableBundle bundle = new PersistableBundle();
+        bundle.putString(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY,
+                Base64.encodeToString(getInferenceInfoBytes(1, 1, 100), Base64.DEFAULT));
+        inferenceInfoStore.addInferenceInfoFromBundle(bundle);
+        List<InferenceInfo> inferenceInfos = inferenceInfoStore.getLatestInferenceInfo(0);
+        assertThat(inferenceInfos).hasSize(1);
+        assertThat(inferenceInfos.get(0).getUid()).isEqualTo(1);
+        assertThat(inferenceInfos.get(0).getStartTimeMs()).isEqualTo(1);
+        assertThat(inferenceInfos.get(0).getEndTimeMs()).isEqualTo(100);
+    }
+
+
+    @Test
+    public void testEvictionAfterMaxAge() throws Exception {
+        PersistableBundle bundle = new PersistableBundle();
+        long testStartTime = System.currentTimeMillis();
+        bundle.putString(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY,
+                Base64.encodeToString(getInferenceInfoBytes(1,  testStartTime - 10,
+                        testStartTime + 100), Base64.DEFAULT));
+        inferenceInfoStore.addInferenceInfoFromBundle(bundle);
+        bundle.putString(OnDeviceSandboxedInferenceService.INFERENCE_INFO_BUNDLE_KEY,
+                Base64.encodeToString(getInferenceInfoBytes(1, testStartTime - 5,
+                        testStartTime + 100), Base64.DEFAULT));
+        inferenceInfoStore.addInferenceInfoFromBundle(bundle);
+        Thread.sleep(1020);
+        List<InferenceInfo> inferenceInfos = inferenceInfoStore.getLatestInferenceInfo(0);
+        assertThat(inferenceInfos).hasSize(2);
+        assertThat(inferenceInfos.get(0).getUid()).isEqualTo(1);
+        assertThat(inferenceInfos.get(0).getStartTimeMs()).isEqualTo(testStartTime - 10);
+        assertThat(inferenceInfos.get(0).getEndTimeMs()).isEqualTo(testStartTime + 100);
+        inferenceInfoStore.addInferenceInfoFromBundle(bundle);
+        List<InferenceInfo> inferenceInfos2 = inferenceInfoStore.getLatestInferenceInfo(0);
+        assertThat(inferenceInfos2).hasSize(1); //previous entries should have been evicted
+    }
+
+    private byte[] getInferenceInfoBytes(int uid, long startTime, long endTime) {
+        com.android.server.ondeviceintelligence.nano.InferenceInfo inferenceInfo =
+                new com.android.server.ondeviceintelligence.nano.InferenceInfo();
+        inferenceInfo.uid = uid;
+        inferenceInfo.startTimeMs = startTime;
+        inferenceInfo.endTimeMs = endTime;
+        return MessageNano.toByteArray(inferenceInfo);
+    }
+}