Pull upstream changes.

Change-Id: I8a1e76e3e42188364ac40c0c51efb1e49292c015

3 files changed, 233 insertions, 133 deletions

diff --git a/java/src/com/google/android/icing/BreakIteratorBatcher.java b/java/src/com/google/android/icing/BreakIteratorBatcher.java
index e3f325b..2b87327 100644
--- a/java/src/com/google/android/icing/BreakIteratorBatcher.java
+++ b/java/src/com/google/android/icing/BreakIteratorBatcher.java
@@ -1,6 +1,19 @@
+// Copyright (C) 2019 Google LLC
+//
+// 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.google.android.icing;
 
-import com.google.android.apps.common.proguard.UsedByNative;
 import java.text.BreakIterator;
 import java.util.ArrayList;
 import java.util.List;
@@ -23,33 +36,27 @@ import java.util.Locale;
  * assertThat(utf16Boundaries).asList().containsExactly(9);
  * }</pre>
  */
-@UsedByNative("jni-cache.cc")
 public class BreakIteratorBatcher {
 
   private final BreakIterator iterator;
 
-  @UsedByNative("jni-cache.cc")
   public BreakIteratorBatcher(Locale locale) {
     this.iterator = BreakIterator.getWordInstance(locale);
   }
 
   /* Direct calls to BreakIterator */
-  @UsedByNative("jni-cache.cc")
   public void setText(String text) {
     iterator.setText(text);
   }
 
-  @UsedByNative("jni-cache.cc")
   public int first() {
     return iterator.first();
   }
 
-  @UsedByNative("jni-cache.cc")
   public int preceding(int utf16Offset) {
     return iterator.preceding(utf16Offset);
   }
 
-  @UsedByNative("jni-cache.cc")
   public int following(int utf16Offset) {
     return iterator.following(utf16Offset);
   }
@@ -60,7 +67,6 @@ public class BreakIteratorBatcher {
    * the end of the text (returns BreakIterator#DONE), then only the results of the previous calls
    * in that batch will be returned.
    */
-  @UsedByNative("jni-cache.cc")
   public int[] next(int batchSize) {
     List<Integer> breakIndices = new ArrayList<>(batchSize);
     for (int i = 0; i < batchSize; ++i) {
diff --git a/java/src/com/google/android/icing/IcingSearchEngine.java b/java/src/com/google/android/icing/IcingSearchEngine.java
index 03a4fbe..ab34f8a 100644
--- a/java/src/com/google/android/icing/IcingSearchEngine.java
+++ b/java/src/com/google/android/icing/IcingSearchEngine.java
@@ -1,28 +1,26 @@
-/*
- * Copyright (C) 2020 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.
- */
+// Copyright (C) 2019 Google LLC
+//
+// 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.google.android.icing;
 
-import android.annotation.NonNull;
 import android.util.Log;
-
+import androidx.annotation.NonNull;
 import com.google.android.icing.proto.DeleteByNamespaceResultProto;
 import com.google.android.icing.proto.DeleteBySchemaTypeResultProto;
 import com.google.android.icing.proto.DeleteResultProto;
 import com.google.android.icing.proto.DocumentProto;
+import com.google.android.icing.proto.GetOptimizeInfoResultProto;
 import com.google.android.icing.proto.GetResultProto;
 import com.google.android.icing.proto.GetSchemaResultProto;
 import com.google.android.icing.proto.GetSchemaTypeResultProto;
@@ -31,6 +29,7 @@ import com.google.android.icing.proto.InitializeResultProto;
 import com.google.android.icing.proto.OptimizeResultProto;
 import com.google.android.icing.proto.PersistToDiskResultProto;
 import com.google.android.icing.proto.PutResultProto;
+import com.google.android.icing.proto.ResetResultProto;
 import com.google.android.icing.proto.ResultSpecProto;
 import com.google.android.icing.proto.SchemaProto;
 import com.google.android.icing.proto.ScoringSpecProto;
@@ -38,6 +37,7 @@ import com.google.android.icing.proto.SearchResultProto;
 import com.google.android.icing.proto.SearchSpecProto;
 import com.google.android.icing.proto.SetSchemaResultProto;
 import com.google.android.icing.proto.StatusProto;
+import com.google.android.icing.protobuf.ExtensionRegistryLite;
 import com.google.android.icing.protobuf.InvalidProtocolBufferException;
 
 /** Java wrapper to access native APIs in external/icing/icing/icing-search-engine.h */
@@ -45,28 +45,25 @@ public final class IcingSearchEngine {
 
   private static final String TAG = "IcingSearchEngine";
 
-  private long mNativePointer;
+  private final long nativePointer;
 
   static {
     // NOTE: This can fail with an UnsatisfiedLinkError
-    System.loadLibrary("icing_jni");
+    System.loadLibrary("icing");
   }
 
-  /**
-   * @throws RuntimeException if IcingSearchEngine fails to be created
-   */
-  public IcingSearchEngine(IcingSearchEngineOptions options) {
-    mNativePointer = nativeCreate(options.toByteArray());
-    if (mNativePointer == 0) {
+  /** @throws IllegalStateException if IcingSearchEngine fails to be created */
+  public IcingSearchEngine(@NonNull IcingSearchEngineOptions options) {
+    nativePointer = nativeCreate(options.toByteArray());
+    if (nativePointer == 0) {
       Log.e(TAG, "Failed to create IcingSearchEngine.");
-      throw new RuntimeException("Failed to create IcingSearchEngine.");
+      throw new IllegalStateException("Failed to create IcingSearchEngine.");
     }
   }
 
-
   @NonNull
   public InitializeResultProto initialize() {
-    byte[] initializeResultBytes = nativeInitialize(mNativePointer);
+    byte[] initializeResultBytes = nativeInitialize(nativePointer);
     if (initializeResultBytes == null) {
       Log.e(TAG, "Received null InitializeResult from native.");
       return InitializeResultProto.newBuilder()
@@ -75,7 +72,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return InitializeResultProto.parseFrom(initializeResultBytes);
+      return InitializeResultProto.parseFrom(
+          initializeResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing InitializeResultProto.", e);
       return InitializeResultProto.newBuilder()
@@ -91,9 +89,9 @@ public final class IcingSearchEngine {
 
   @NonNull
   public SetSchemaResultProto setSchema(
-          @NonNull SchemaProto schema, boolean ignoreErrorsAndDeleteDocuments) {
+      @NonNull SchemaProto schema, boolean ignoreErrorsAndDeleteDocuments) {
     byte[] setSchemaResultBytes =
-        nativeSetSchema(mNativePointer, schema.toByteArray(), ignoreErrorsAndDeleteDocuments);
+        nativeSetSchema(nativePointer, schema.toByteArray(), ignoreErrorsAndDeleteDocuments);
     if (setSchemaResultBytes == null) {
       Log.e(TAG, "Received null SetSchemaResultProto from native.");
       return SetSchemaResultProto.newBuilder()
@@ -102,7 +100,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return SetSchemaResultProto.parseFrom(setSchemaResultBytes);
+      return SetSchemaResultProto.parseFrom(
+          setSchemaResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing SetSchemaResultProto.", e);
       return SetSchemaResultProto.newBuilder()
@@ -113,7 +112,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public GetSchemaResultProto getSchema() {
-    byte[] getSchemaResultBytes = nativeGetSchema(mNativePointer);
+    byte[] getSchemaResultBytes = nativeGetSchema(nativePointer);
     if (getSchemaResultBytes == null) {
       Log.e(TAG, "Received null GetSchemaResultProto from native.");
       return GetSchemaResultProto.newBuilder()
@@ -122,7 +121,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return GetSchemaResultProto.parseFrom(getSchemaResultBytes);
+      return GetSchemaResultProto.parseFrom(
+          getSchemaResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing GetSchemaResultProto.", e);
       return GetSchemaResultProto.newBuilder()
@@ -133,7 +133,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public GetSchemaTypeResultProto getSchemaType(@NonNull String schemaType) {
-    byte[] getSchemaTypeResultBytes = nativeGetSchemaType(mNativePointer, schemaType);
+    byte[] getSchemaTypeResultBytes = nativeGetSchemaType(nativePointer, schemaType);
     if (getSchemaTypeResultBytes == null) {
       Log.e(TAG, "Received null GetSchemaTypeResultProto from native.");
       return GetSchemaTypeResultProto.newBuilder()
@@ -142,7 +142,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return GetSchemaTypeResultProto.parseFrom(getSchemaTypeResultBytes);
+      return GetSchemaTypeResultProto.parseFrom(
+          getSchemaTypeResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing GetSchemaTypeResultProto.", e);
       return GetSchemaTypeResultProto.newBuilder()
@@ -153,7 +154,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public PutResultProto put(@NonNull DocumentProto document) {
-    byte[] putResultBytes = nativePut(mNativePointer, document.toByteArray());
+    byte[] putResultBytes = nativePut(nativePointer, document.toByteArray());
     if (putResultBytes == null) {
       Log.e(TAG, "Received null PutResultProto from native.");
       return PutResultProto.newBuilder()
@@ -162,7 +163,7 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return PutResultProto.parseFrom(putResultBytes);
+      return PutResultProto.parseFrom(putResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing PutResultProto.", e);
       return PutResultProto.newBuilder()
@@ -173,7 +174,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public GetResultProto get(@NonNull String namespace, @NonNull String uri) {
-    byte[] getResultBytes = nativeGet(mNativePointer, namespace, uri);
+    byte[] getResultBytes = nativeGet(nativePointer, namespace, uri);
     if (getResultBytes == null) {
       Log.e(TAG, "Received null GetResultProto from native.");
       return GetResultProto.newBuilder()
@@ -182,7 +183,7 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return GetResultProto.parseFrom(getResultBytes);
+      return GetResultProto.parseFrom(getResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing GetResultProto.", e);
       return GetResultProto.newBuilder()
@@ -193,10 +194,12 @@ public final class IcingSearchEngine {
 
   @NonNull
   public SearchResultProto search(
-          @NonNull SearchSpecProto searchSpec, @NonNull ScoringSpecProto scoringSpec, @NonNull ResultSpecProto resultSpec) {
+      @NonNull SearchSpecProto searchSpec,
+      @NonNull ScoringSpecProto scoringSpec,
+      @NonNull ResultSpecProto resultSpec) {
     byte[] searchResultBytes =
         nativeSearch(
-            mNativePointer,
+            nativePointer,
             searchSpec.toByteArray(),
             scoringSpec.toByteArray(),
             resultSpec.toByteArray());
@@ -208,7 +211,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return SearchResultProto.parseFrom(searchResultBytes);
+      return SearchResultProto.parseFrom(
+          searchResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing SearchResultProto.", e);
       return SearchResultProto.newBuilder()
@@ -219,7 +223,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public DeleteResultProto delete(@NonNull String namespace, @NonNull String uri) {
-    byte[] deleteResultBytes = nativeDelete(mNativePointer, namespace, uri);
+    byte[] deleteResultBytes = nativeDelete(nativePointer, namespace, uri);
     if (deleteResultBytes == null) {
       Log.e(TAG, "Received null DeleteResultProto from native.");
       return DeleteResultProto.newBuilder()
@@ -228,7 +232,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return DeleteResultProto.parseFrom(deleteResultBytes);
+      return DeleteResultProto.parseFrom(
+          deleteResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing DeleteResultProto.", e);
       return DeleteResultProto.newBuilder()
@@ -239,7 +244,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public DeleteByNamespaceResultProto deleteByNamespace(@NonNull String namespace) {
-    byte[] deleteByNamespaceResultBytes = nativeDeleteByNamespace(mNativePointer, namespace);
+    byte[] deleteByNamespaceResultBytes = nativeDeleteByNamespace(nativePointer, namespace);
     if (deleteByNamespaceResultBytes == null) {
       Log.e(TAG, "Received null DeleteByNamespaceResultProto from native.");
       return DeleteByNamespaceResultProto.newBuilder()
@@ -248,7 +253,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return DeleteByNamespaceResultProto.parseFrom(deleteByNamespaceResultBytes);
+      return DeleteByNamespaceResultProto.parseFrom(
+          deleteByNamespaceResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing DeleteByNamespaceResultProto.", e);
       return DeleteByNamespaceResultProto.newBuilder()
@@ -259,7 +265,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public DeleteBySchemaTypeResultProto deleteBySchemaType(@NonNull String schemaType) {
-    byte[] deleteBySchemaTypeResultBytes = nativeDeleteBySchemaType(mNativePointer, schemaType);
+    byte[] deleteBySchemaTypeResultBytes = nativeDeleteBySchemaType(nativePointer, schemaType);
     if (deleteBySchemaTypeResultBytes == null) {
       Log.e(TAG, "Received null DeleteBySchemaTypeResultProto from native.");
       return DeleteBySchemaTypeResultProto.newBuilder()
@@ -268,7 +274,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return DeleteBySchemaTypeResultProto.parseFrom(deleteBySchemaTypeResultBytes);
+      return DeleteBySchemaTypeResultProto.parseFrom(
+          deleteBySchemaTypeResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing DeleteBySchemaTypeResultProto.", e);
       return DeleteBySchemaTypeResultProto.newBuilder()
@@ -279,7 +286,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public PersistToDiskResultProto persistToDisk() {
-    byte[] persistToDiskResultBytes = nativePersistToDisk(mNativePointer);
+    byte[] persistToDiskResultBytes = nativePersistToDisk(nativePointer);
     if (persistToDiskResultBytes == null) {
       Log.e(TAG, "Received null PersistToDiskResultProto from native.");
       return PersistToDiskResultProto.newBuilder()
@@ -288,7 +295,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return PersistToDiskResultProto.parseFrom(persistToDiskResultBytes);
+      return PersistToDiskResultProto.parseFrom(
+          persistToDiskResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing PersistToDiskResultProto.", e);
       return PersistToDiskResultProto.newBuilder()
@@ -299,7 +307,7 @@ public final class IcingSearchEngine {
 
   @NonNull
   public OptimizeResultProto optimize() {
-    byte[] optimizeResultBytes = nativeOptimize(mNativePointer);
+    byte[] optimizeResultBytes = nativeOptimize(nativePointer);
     if (optimizeResultBytes == null) {
       Log.e(TAG, "Received null OptimizeResultProto from native.");
       return OptimizeResultProto.newBuilder()
@@ -308,7 +316,8 @@ public final class IcingSearchEngine {
     }
 
     try {
-      return OptimizeResultProto.parseFrom(optimizeResultBytes);
+      return OptimizeResultProto.parseFrom(
+          optimizeResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
     } catch (InvalidProtocolBufferException e) {
       Log.e(TAG, "Error parsing OptimizeResultProto.", e);
       return OptimizeResultProto.newBuilder()
@@ -317,31 +326,77 @@ public final class IcingSearchEngine {
     }
   }
 
+  @NonNull
+  public GetOptimizeInfoResultProto getOptimizeInfo() {
+    byte[] getOptimizeInfoResultBytes = nativeGetOptimizeInfo(nativePointer);
+    if (getOptimizeInfoResultBytes == null) {
+      Log.e(TAG, "Received null GetOptimizeInfoResultProto from native.");
+      return GetOptimizeInfoResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
+    }
+
+    try {
+      return GetOptimizeInfoResultProto.parseFrom(
+          getOptimizeInfoResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
+    } catch (InvalidProtocolBufferException e) {
+      Log.e(TAG, "Error parsing GetOptimizeInfoResultProto.", e);
+      return GetOptimizeInfoResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
+    }
+  }
+
+  @NonNull
+  public ResetResultProto reset() {
+    byte[] resetResultBytes = nativeReset(nativePointer);
+    if (resetResultBytes == null) {
+      Log.e(TAG, "Received null ResetResultProto from native.");
+      return ResetResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
+    }
+
+    try {
+      return ResetResultProto.parseFrom(
+          resetResultBytes, ExtensionRegistryLite.getGeneratedRegistry());
+    } catch (InvalidProtocolBufferException e) {
+      Log.e(TAG, "Error parsing ResetResultProto.", e);
+      return ResetResultProto.newBuilder()
+          .setStatus(StatusProto.newBuilder().setCode(StatusProto.Code.INTERNAL))
+          .build();
+    }
+  }
+
   private static native long nativeCreate(byte[] icingSearchEngineOptionsBytes);
 
-  private static native byte[] nativeInitialize(long mNativePointer);
+  private static native byte[] nativeInitialize(long nativePointer);
 
   private static native byte[] nativeSetSchema(
-      long mNativePointer, byte[] schemaBytes, boolean ignoreErrorsAndDeleteDocuments);
+      long nativePointer, byte[] schemaBytes, boolean ignoreErrorsAndDeleteDocuments);
 
-  private static native byte[] nativeGetSchema(long mNativePointer);
+  private static native byte[] nativeGetSchema(long nativePointer);
 
-  private static native byte[] nativeGetSchemaType(long mNativePointer, String schemaType);
+  private static native byte[] nativeGetSchemaType(long nativePointer, String schemaType);
 
-  private static native byte[] nativePut(long mNativePointer, byte[] documentBytes);
+  private static native byte[] nativePut(long nativePointer, byte[] documentBytes);
 
-  private static native byte[] nativeGet(long mNativePointer, String namespace, String uri);
+  private static native byte[] nativeGet(long nativePointer, String namespace, String uri);
 
   private static native byte[] nativeSearch(
-      long mNativePointer, byte[] searchSpecBytes, byte[] scoringSpecBytes, byte[] resultSpecBytes);
+      long nativePointer, byte[] searchSpecBytes, byte[] scoringSpecBytes, byte[] resultSpecBytes);
+
+  private static native byte[] nativeDelete(long nativePointer, String namespace, String uri);
+
+  private static native byte[] nativeDeleteByNamespace(long nativePointer, String namespace);
 
-  private static native byte[] nativeDelete(long mNativePointer, String namespace, String uri);
+  private static native byte[] nativeDeleteBySchemaType(long nativePointer, String schemaType);
 
-  private static native byte[] nativeDeleteByNamespace(long mNativePointer, String namespace);
+  private static native byte[] nativePersistToDisk(long nativePointer);
 
-  private static native byte[] nativeDeleteBySchemaType(long mNativePointer, String schemaType);
+  private static native byte[] nativeOptimize(long nativePointer);
 
-  private static native byte[] nativePersistToDisk(long mNativePointer);
+  private static native byte[] nativeGetOptimizeInfo(long nativePointer);
 
-  private static native byte[] nativeOptimize(long mNativePointer);
+  private static native byte[] nativeReset(long nativePointer);
 }
diff --git a/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java b/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java
index 5c502d9..3c7c2ea 100644
--- a/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java
+++ b/java/tests/instrumentation/src/com/google/android/icing/IcingSearchEngineTest.java
@@ -1,31 +1,27 @@
-/*
- * Copyright (C) 2020 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.
- */
+// Copyright (C) 2019 Google LLC
+//
+// 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.google.android.icing;
 
 import static com.google.common.truth.Truth.assertThat;
 
-import android.content.Context;
-
-import androidx.test.InstrumentationRegistry;
-
+import androidx.test.core.app.ApplicationProvider;
 import com.google.android.icing.proto.DeleteByNamespaceResultProto;
 import com.google.android.icing.proto.DeleteBySchemaTypeResultProto;
 import com.google.android.icing.proto.DeleteResultProto;
 import com.google.android.icing.proto.DocumentProto;
+import com.google.android.icing.proto.GetOptimizeInfoResultProto;
 import com.google.android.icing.proto.GetResultProto;
 import com.google.android.icing.proto.GetSchemaResultProto;
 import com.google.android.icing.proto.GetSchemaTypeResultProto;
@@ -38,6 +34,7 @@ import com.google.android.icing.proto.PersistToDiskResultProto;
 import com.google.android.icing.proto.PropertyConfigProto;
 import com.google.android.icing.proto.PropertyProto;
 import com.google.android.icing.proto.PutResultProto;
+import com.google.android.icing.proto.ResetResultProto;
 import com.google.android.icing.proto.ResultSpecProto;
 import com.google.android.icing.proto.SchemaProto;
 import com.google.android.icing.proto.SchemaTypeConfigProto;
@@ -47,25 +44,24 @@ import com.google.android.icing.proto.SearchSpecProto;
 import com.google.android.icing.proto.SetSchemaResultProto;
 import com.google.android.icing.proto.StatusProto;
 import com.google.android.icing.proto.TermMatchType;
-
+import com.google.android.libraries.mdi.search.IcingSearchEngine;
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
-import org.junit.runners.JUnit4;
+import org.robolectric.RobolectricTestRunner;
 
 /**
  * This test is not intended to fully test the functionality of each API. But rather to test the JNI
- * wrapper of Icing library.
+ * wrapper and Java interfaces of Icing library {@link IcingSearchEngine}.
  */
-@RunWith(JUnit4.class)
+@RunWith(RobolectricTestRunner.class)
 public final class IcingSearchEngineTest {
 
   private static final String EMAIL_TYPE = "Email";
 
-  private Context mContext;
-  private String mFilesDir;
+  private String filesDir;
 
-  static SchemaTypeConfigProto createEmailTypeConfig() {
+  private static SchemaTypeConfigProto createEmailTypeConfig() {
     return SchemaTypeConfigProto.newBuilder()
         .setSchemaType(EMAIL_TYPE)
         .addProperties(
@@ -89,7 +85,7 @@ public final class IcingSearchEngineTest {
         .build();
   }
 
-  static DocumentProto createEmailDocument(String namespace, String uri) {
+  private static DocumentProto createEmailDocument(String namespace, String uri) {
     return DocumentProto.newBuilder()
         .setNamespace(namespace)
         .setUri(uri)
@@ -100,14 +96,13 @@ public final class IcingSearchEngineTest {
 
   @Before
   public void setUp() throws Exception {
-    mContext = InstrumentationRegistry.getInstrumentation().getContext();
-    mFilesDir = mContext.getFilesDir().getCanonicalPath();
+    filesDir = ApplicationProvider.getApplicationContext().getFilesDir().getCanonicalPath();
   }
 
   @Test
   public void testInitialize() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
 
     InitializeResultProto initializeResultProto = icing.initialize();
@@ -117,9 +112,9 @@ public final class IcingSearchEngineTest {
   @Test
   public void testSetAndGetSchema() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
-    InitializeResultProto initializeResultProto = icing.initialize();
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
     SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
@@ -140,14 +135,18 @@ public final class IcingSearchEngineTest {
   @Test
   public void testPutAndGetDocuments() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
-    InitializeResultProto initializeResultProto = icing.initialize();
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
     SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
-    SetSchemaResultProto setSchemaResultProto =
-        icing.setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false);
+    assertThat(
+            icing
+                .setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
 
     DocumentProto emailDocument = createEmailDocument("namespace", "uri");
     PutResultProto putResultProto = icing.put(emailDocument);
@@ -161,20 +160,24 @@ public final class IcingSearchEngineTest {
   @Test
   public void testSearch() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
-    InitializeResultProto initializeResultProto = icing.initialize();
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
     SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
-    SetSchemaResultProto setSchemaResultProto =
-        icing.setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false);
+    assertThat(
+            icing
+                .setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
 
     DocumentProto emailDocument =
         createEmailDocument("namespace", "uri").toBuilder()
             .addProperties(PropertyProto.newBuilder().setName("subject").addStringValues("foo"))
             .build();
-    PutResultProto putResultProto = icing.put(emailDocument);
+    assertThat(icing.put(emailDocument).getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     SearchSpecProto searchSpec =
         SearchSpecProto.newBuilder()
@@ -195,17 +198,21 @@ public final class IcingSearchEngineTest {
   @Test
   public void testDelete() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
-    InitializeResultProto initializeResultProto = icing.initialize();
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
     SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
-    SetSchemaResultProto setSchemaResultProto =
-        icing.setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false);
+    assertThat(
+            icing
+                .setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
 
     DocumentProto emailDocument = createEmailDocument("namespace", "uri");
-    PutResultProto putResultProto = icing.put(emailDocument);
+    assertThat(icing.put(emailDocument).getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     DeleteResultProto deleteResultProto = icing.delete("namespace", "uri");
     assertThat(deleteResultProto.getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
@@ -217,17 +224,21 @@ public final class IcingSearchEngineTest {
   @Test
   public void testDeleteByNamespace() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
-    InitializeResultProto initializeResultProto = icing.initialize();
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
     SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
-    SetSchemaResultProto setSchemaResultProto =
-        icing.setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false);
+    assertThat(
+            icing
+                .setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
 
     DocumentProto emailDocument = createEmailDocument("namespace", "uri");
-    PutResultProto putResultProto = icing.put(emailDocument);
+    assertThat(icing.put(emailDocument).getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     DeleteByNamespaceResultProto deleteByNamespaceResultProto =
         icing.deleteByNamespace("namespace");
@@ -240,17 +251,21 @@ public final class IcingSearchEngineTest {
   @Test
   public void testDeleteBySchemaType() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
-    InitializeResultProto initializeResultProto = icing.initialize();
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     SchemaTypeConfigProto emailTypeConfig = createEmailTypeConfig();
     SchemaProto schema = SchemaProto.newBuilder().addTypes(emailTypeConfig).build();
-    SetSchemaResultProto setSchemaResultProto =
-        icing.setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false);
+    assertThat(
+            icing
+                .setSchema(schema, /*ignoreErrorsAndDeleteDocuments=*/ false)
+                .getStatus()
+                .getCode())
+        .isEqualTo(StatusProto.Code.OK);
 
     DocumentProto emailDocument = createEmailDocument("namespace", "uri");
-    PutResultProto putResultProto = icing.put(emailDocument);
+    assertThat(icing.put(emailDocument).getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     DeleteBySchemaTypeResultProto deleteBySchemaTypeResultProto =
         icing.deleteBySchemaType(EMAIL_TYPE);
@@ -263,9 +278,9 @@ public final class IcingSearchEngineTest {
   @Test
   public void testPersistToDisk() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
-    InitializeResultProto initializeResultProto = icing.initialize();
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     PersistToDiskResultProto persistToDiskResultProto = icing.persistToDisk();
     assertThat(persistToDiskResultProto.getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
@@ -274,11 +289,35 @@ public final class IcingSearchEngineTest {
   @Test
   public void testOptimize() throws Exception {
     IcingSearchEngineOptions options =
-        IcingSearchEngineOptions.newBuilder().setBaseDir(mFilesDir).build();
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
     IcingSearchEngine icing = new IcingSearchEngine(options);
-    InitializeResultProto initializeResultProto = icing.initialize();
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
 
     OptimizeResultProto optimizeResultProto = icing.optimize();
     assertThat(optimizeResultProto.getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
   }
+
+  @Test
+  public void testGetOptimizeInfo() throws Exception {
+    IcingSearchEngineOptions options =
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
+    IcingSearchEngine icing = new IcingSearchEngine(options);
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
+
+    GetOptimizeInfoResultProto getOptimizeInfoResultProto = icing.getOptimizeInfo();
+    assertThat(getOptimizeInfoResultProto.getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
+    assertThat(getOptimizeInfoResultProto.getOptimizableDocs()).isEqualTo(0);
+    assertThat(getOptimizeInfoResultProto.getEstimatedOptimizableBytes()).isEqualTo(0);
+  }
+
+  @Test
+  public void testReset() throws Exception {
+    IcingSearchEngineOptions options =
+        IcingSearchEngineOptions.newBuilder().setBaseDir(filesDir).build();
+    IcingSearchEngine icing = new IcingSearchEngine(options);
+    assertThat(icing.initialize().getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
+
+    ResetResultProto resetResultProto = icing.reset();
+    assertThat(resetResultProto.getStatus().getCode()).isEqualTo(StatusProto.Code.OK);
+  }
 }