Abstract FutureAppSearchSession

Context: This is to make writing the test much easier for other dependent classes like the syncer.
Flag: android.app.appfunctions.flags.enable_app_function_manager
Test: atest FrameworksAppFunctionsTests -c
Bug: 357551503
Change-Id: I853ab02013c98d395ffd1e02f1f163d0677b4fd4

5 files changed, 257 insertions, 200 deletions

diff --git a/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java b/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
index 39aa27ae794e..b1c25c4f3c61 100644
--- a/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
+++ b/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSession.java
@@ -18,11 +18,8 @@ package com.android.server.appfunctions;
 
 import android.annotation.NonNull;
 import android.app.appsearch.AppSearchBatchResult;
-import android.app.appsearch.AppSearchManager;
-import android.app.appsearch.AppSearchManager.SearchContext;
 import android.app.appsearch.AppSearchResult;
 import android.app.appsearch.AppSearchSession;
-import android.app.appsearch.BatchResultCallback;
 import android.app.appsearch.GenericDocument;
 import android.app.appsearch.GetByDocumentIdRequest;
 import android.app.appsearch.GetSchemaResponse;
@@ -33,7 +30,6 @@ import android.app.appsearch.SearchResults;
 import android.app.appsearch.SearchSpec;
 import android.app.appsearch.SetSchemaRequest;
 import android.app.appsearch.SetSchemaResponse;
-import android.util.Slog;
 
 import com.android.internal.infra.AndroidFuture;
 
@@ -44,28 +40,11 @@ import java.util.Objects;
 import java.util.concurrent.Executor;
 
 /** A future API wrapper of {@link AppSearchSession} APIs. */
-public class FutureAppSearchSession implements Closeable {
-    private static final String TAG = FutureAppSearchSession.class.getSimpleName();
-    private final Executor mExecutor;
-    private final AndroidFuture<AppSearchResult<AppSearchSession>> mSettableSessionFuture;
-
-    public FutureAppSearchSession(
-            @NonNull AppSearchManager appSearchManager,
-            @NonNull Executor executor,
-            @NonNull SearchContext appSearchContext) {
-        Objects.requireNonNull(appSearchManager);
-        Objects.requireNonNull(executor);
-        Objects.requireNonNull(appSearchContext);
-
-        mExecutor = executor;
-        mSettableSessionFuture = new AndroidFuture<>();
-        appSearchManager.createSearchSession(
-                appSearchContext, mExecutor, mSettableSessionFuture::complete);
-    }
+public interface FutureAppSearchSession extends Closeable {
 
     /** Converts a failed app search result codes into an exception. */
     @NonNull
-    public static Exception failedResultToException(@NonNull AppSearchResult<?> appSearchResult) {
+    static Exception failedResultToException(@NonNull AppSearchResult<?> appSearchResult) {
         return switch (appSearchResult.getResultCode()) {
             case AppSearchResult.RESULT_INVALID_ARGUMENT ->
                     new IllegalArgumentException(appSearchResult.getErrorMessage());
@@ -77,114 +56,39 @@ public class FutureAppSearchSession implements Closeable {
         };
     }
 
-    private AndroidFuture<AppSearchSession> getSessionAsync() {
-        return mSettableSessionFuture.thenApply(
-                result -> {
-                    if (result.isSuccess()) {
-                        return result.getResultValue();
-                    } else {
-                        throw new RuntimeException(failedResultToException(result));
-                    }
-                });
-    }
-
-    /** Gets the schema for a given app search session. */
-    public AndroidFuture<GetSchemaResponse> getSchema() {
-        return getSessionAsync()
-                .thenCompose(
-                        session -> {
-                            AndroidFuture<AppSearchResult<GetSchemaResponse>>
-                                    settableSchemaResponse = new AndroidFuture<>();
-                            session.getSchema(mExecutor, settableSchemaResponse::complete);
-                            return settableSchemaResponse.thenApply(
-                                    result -> {
-                                        if (result.isSuccess()) {
-                                            return result.getResultValue();
-                                        } else {
-                                            throw new RuntimeException(
-                                                    failedResultToException(result));
-                                        }
-                                    });
-                        });
-    }
-
-    /** Sets the schema for a given app search session. */
-    public AndroidFuture<SetSchemaResponse> setSchema(@NonNull SetSchemaRequest setSchemaRequest) {
-        return getSessionAsync()
-                .thenCompose(
-                        session -> {
-                            AndroidFuture<AppSearchResult<SetSchemaResponse>>
-                                    settableSchemaResponse = new AndroidFuture<>();
-                            session.setSchema(
-                                    setSchemaRequest,
-                                    mExecutor,
-                                    mExecutor,
-                                    settableSchemaResponse::complete);
-                            return settableSchemaResponse.thenApply(
-                                    result -> {
-                                        if (result.isSuccess()) {
-                                            return result.getResultValue();
-                                        } else {
-                                            throw new RuntimeException(
-                                                    failedResultToException(result));
-                                        }
-                                    });
-                        });
-    }
+    /**
+     * Sets the schema that represents the organizational structure of data within the AppSearch
+     * database.
+     */
+    AndroidFuture<SetSchemaResponse> setSchema(@NonNull SetSchemaRequest setSchemaRequest);
 
-    /** Indexes documents into the AppSearchSession database. */
-    public AndroidFuture<AppSearchBatchResult<String, Void>> put(
-            @NonNull PutDocumentsRequest putDocumentsRequest) {
-        return getSessionAsync()
-                .thenCompose(
-                        session -> {
-                            AndroidFuture<AppSearchBatchResult<String, Void>> batchResultFuture =
-                                    new AndroidFuture<>();
+    /** Retrieves the schema most recently successfully provided to {@code setSchema}. */
+    AndroidFuture<GetSchemaResponse> getSchema();
 
-                            session.put(
-                                    putDocumentsRequest, mExecutor, batchResultFuture::complete);
-                            return batchResultFuture;
-                        });
-    }
+    /** Indexes documents into the {@link AppSearchSession} database. */
+    AndroidFuture<AppSearchBatchResult<String, Void>> put(
+            @NonNull PutDocumentsRequest putDocumentsRequest);
 
-    /** Removes documents from the AppSearchSession database. */
-    public AndroidFuture<AppSearchBatchResult<String, Void>> remove(
-            @NonNull RemoveByDocumentIdRequest removeRequest) {
-        return getSessionAsync()
-                .thenCompose(
-                        session -> {
-                            AndroidFuture<AppSearchBatchResult<String, Void>>
-                                    settableBatchResultFuture = new AndroidFuture<>();
-                            session.remove(
-                                    removeRequest,
-                                    mExecutor,
-                                    new BatchResultCallbackAdapter<>(settableBatchResultFuture));
-                            return settableBatchResultFuture;
-                        });
-    }
+    /** Removes {@link GenericDocument} from the index by Query. */
+    AndroidFuture<AppSearchBatchResult<String, Void>> remove(
+            @NonNull RemoveByDocumentIdRequest removeRequest);
 
     /**
-     * Retrieves documents from the open AppSearchSession that match a given query string and type
-     * of search provided.
+     * Gets {@link GenericDocument} objects by document IDs in a namespace from the {@link
+     * AppSearchSession} database.
      */
-    public AndroidFuture<FutureSearchResults> search(
-            @NonNull String queryExpression, @NonNull SearchSpec searchSpec) {
-        return getSessionAsync()
-                .thenApply(session -> session.search(queryExpression, searchSpec))
-                .thenApply(result -> new FutureSearchResults(result, mExecutor));
-    }
+    AndroidFuture<AppSearchBatchResult<String, GenericDocument>> getByDocumentId(
+            GetByDocumentIdRequest getRequest);
 
-    @Override
-    public void close() throws IOException {
-        try {
-            getSessionAsync().get().close();
-        } catch (Exception ex) {
-            Slog.e(TAG, "Failed to close app search session", ex);
-        }
-    }
+    /**
+     * Retrieves documents from the open {@link AppSearchSession} that match a given query string
+     * and type of search provided.
+     */
+    AndroidFuture<FutureSearchResults> search(
+            @NonNull String queryExpression, @NonNull SearchSpec searchSpec);
 
     /** A future API wrapper of {@link android.app.appsearch.SearchResults}. */
-    public static class FutureSearchResults {
+    class FutureSearchResults {
         private final SearchResults mSearchResults;
         private final Executor mExecutor;
 
@@ -209,56 +113,4 @@ public class FutureAppSearchSession implements Closeable {
                     });
         }
     }
-
-    /** A future API to retrieve a document by its id from the local AppSearch session. */
-    public AndroidFuture<GenericDocument> getByDocumentId(
-            @NonNull String documentId, @NonNull String namespace) {
-        Objects.requireNonNull(documentId);
-        Objects.requireNonNull(namespace);
-
-        GetByDocumentIdRequest request =
-                new GetByDocumentIdRequest.Builder(namespace).addIds(documentId).build();
-        return getSessionAsync()
-                .thenCompose(
-                        session -> {
-                            AndroidFuture<AppSearchBatchResult<String, GenericDocument>>
-                                    batchResultFuture = new AndroidFuture<>();
-                            session.getByDocumentId(
-                                    request,
-                                    mExecutor,
-                                    new BatchResultCallbackAdapter<>(batchResultFuture));
-
-                            return batchResultFuture.thenApply(
-                                    batchResult ->
-                                            getGenericDocumentFromBatchResult(
-                                                    batchResult, documentId));
-                        });
-    }
-
-    private static GenericDocument getGenericDocumentFromBatchResult(
-            AppSearchBatchResult<String, GenericDocument> result, String documentId) {
-        if (result.isSuccess()) {
-            return result.getSuccesses().get(documentId);
-        }
-        throw new IllegalArgumentException("No document in the result for id: " + documentId);
-    }
-
-    private static final class BatchResultCallbackAdapter<K, V>
-            implements BatchResultCallback<K, V> {
-        private final AndroidFuture<AppSearchBatchResult<K, V>> mFuture;
-
-        BatchResultCallbackAdapter(AndroidFuture<AppSearchBatchResult<K, V>> future) {
-            mFuture = future;
-        }
-
-        @Override
-        public void onResult(@NonNull AppSearchBatchResult<K, V> result) {
-            mFuture.complete(result);
-        }
-
-        @Override
-        public void onSystemError(Throwable t) {
-            mFuture.completeExceptionally(t);
-        }
-    }
 }
diff --git a/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSessionImpl.java b/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSessionImpl.java
new file mode 100644
index 000000000000..e78f390a4825
--- /dev/null
+++ b/services/appfunctions/java/com/android/server/appfunctions/FutureAppSearchSessionImpl.java
@@ -0,0 +1,203 @@
+/*
+ * 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.appfunctions;
+
+import static com.android.server.appfunctions.FutureAppSearchSession.failedResultToException;
+
+import android.annotation.NonNull;
+import android.app.appsearch.AppSearchBatchResult;
+import android.app.appsearch.AppSearchManager;
+import android.app.appsearch.AppSearchManager.SearchContext;
+import android.app.appsearch.AppSearchResult;
+import android.app.appsearch.AppSearchSession;
+import android.app.appsearch.BatchResultCallback;
+import android.app.appsearch.GenericDocument;
+import android.app.appsearch.GetByDocumentIdRequest;
+import android.app.appsearch.GetSchemaResponse;
+import android.app.appsearch.PutDocumentsRequest;
+import android.app.appsearch.RemoveByDocumentIdRequest;
+import android.app.appsearch.SearchSpec;
+import android.app.appsearch.SetSchemaRequest;
+import android.app.appsearch.SetSchemaResponse;
+
+import com.android.internal.infra.AndroidFuture;
+
+import java.io.IOException;
+import java.util.Objects;
+import java.util.concurrent.Executor;
+
+/** Implementation of {@link FutureAppSearchSession} */
+public class FutureAppSearchSessionImpl implements FutureAppSearchSession {
+
+    private static final String TAG = FutureAppSearchSession.class.getSimpleName();
+    private final Executor mExecutor;
+    private final AndroidFuture<AppSearchResult<AppSearchSession>> mSettableSessionFuture;
+
+    public FutureAppSearchSessionImpl(
+            @NonNull AppSearchManager appSearchManager,
+            @NonNull Executor executor,
+            @NonNull SearchContext appSearchContext) {
+        Objects.requireNonNull(appSearchManager);
+        Objects.requireNonNull(executor);
+        Objects.requireNonNull(appSearchContext);
+
+        mExecutor = executor;
+        mSettableSessionFuture = new AndroidFuture<>();
+        appSearchManager.createSearchSession(
+                appSearchContext, mExecutor, mSettableSessionFuture::complete);
+    }
+
+    private AndroidFuture<AppSearchSession> getSessionAsync() {
+        return mSettableSessionFuture.thenApply(
+                result -> {
+                    if (result.isSuccess()) {
+                        return result.getResultValue();
+                    } else {
+                        throw new RuntimeException(failedResultToException(result));
+                    }
+                });
+    }
+
+    @Override
+    public AndroidFuture<SetSchemaResponse> setSchema(@NonNull SetSchemaRequest setSchemaRequest) {
+        Objects.requireNonNull(setSchemaRequest);
+
+        return getSessionAsync()
+                .thenCompose(
+                        session -> {
+                            AndroidFuture<AppSearchResult<SetSchemaResponse>>
+                                    settableSchemaResponse = new AndroidFuture<>();
+                            session.setSchema(
+                                    setSchemaRequest,
+                                    mExecutor,
+                                    mExecutor,
+                                    settableSchemaResponse::complete);
+                            return settableSchemaResponse.thenApply(
+                                    result -> {
+                                        if (result.isSuccess()) {
+                                            return result.getResultValue();
+                                        } else {
+                                            throw new RuntimeException(
+                                                    failedResultToException(result));
+                                        }
+                                    });
+                        });
+    }
+
+    @Override
+    public AndroidFuture<GetSchemaResponse> getSchema() {
+        return getSessionAsync()
+                .thenCompose(
+                        session -> {
+                            AndroidFuture<AppSearchResult<GetSchemaResponse>>
+                                    settableSchemaResponse = new AndroidFuture<>();
+                            session.getSchema(mExecutor, settableSchemaResponse::complete);
+                            return settableSchemaResponse.thenApply(
+                                    result -> {
+                                        if (result.isSuccess()) {
+                                            return result.getResultValue();
+                                        } else {
+                                            throw new RuntimeException(
+                                                    failedResultToException(result));
+                                        }
+                                    });
+                        });
+    }
+
+    @Override
+    public AndroidFuture<AppSearchBatchResult<String, Void>> put(
+            @NonNull PutDocumentsRequest putDocumentsRequest) {
+        Objects.requireNonNull(putDocumentsRequest);
+
+        return getSessionAsync()
+                .thenCompose(
+                        session -> {
+                            AndroidFuture<AppSearchBatchResult<String, Void>> batchResultFuture =
+                                    new AndroidFuture<>();
+
+                            session.put(
+                                    putDocumentsRequest, mExecutor, batchResultFuture::complete);
+                            return batchResultFuture;
+                        });
+    }
+
+    @Override
+    public AndroidFuture<AppSearchBatchResult<String, Void>> remove(
+            @NonNull RemoveByDocumentIdRequest removeRequest) {
+        Objects.requireNonNull(removeRequest);
+
+        return getSessionAsync()
+                .thenCompose(
+                        session -> {
+                            AndroidFuture<AppSearchBatchResult<String, Void>>
+                                    settableBatchResultFuture = new AndroidFuture<>();
+                            session.remove(
+                                    removeRequest,
+                                    mExecutor,
+                                    new BatchResultCallbackAdapter<>(settableBatchResultFuture));
+                            return settableBatchResultFuture;
+                        });
+    }
+
+    @Override
+    public AndroidFuture<AppSearchBatchResult<String, GenericDocument>> getByDocumentId(
+            @NonNull GetByDocumentIdRequest getRequest) {
+        Objects.requireNonNull(getRequest);
+
+        return getSessionAsync()
+                .thenCompose(
+                        session -> {
+                            AndroidFuture<AppSearchBatchResult<String, GenericDocument>>
+                                    batchResultFuture = new AndroidFuture<>();
+                            session.getByDocumentId(
+                                    getRequest,
+                                    mExecutor,
+                                    new BatchResultCallbackAdapter<>(batchResultFuture));
+                            return batchResultFuture;
+                        });
+    }
+
+    @Override
+    public AndroidFuture<FutureSearchResults> search(
+            @NonNull String queryExpression, @NonNull SearchSpec searchSpec) {
+        return getSessionAsync()
+                .thenApply(session -> session.search(queryExpression, searchSpec))
+                .thenApply(result -> new FutureSearchResults(result, mExecutor));
+    }
+
+    @Override
+    public void close() throws IOException {}
+
+    private static final class BatchResultCallbackAdapter<K, V>
+            implements BatchResultCallback<K, V> {
+        private final AndroidFuture<AppSearchBatchResult<K, V>> mFuture;
+
+        BatchResultCallbackAdapter(AndroidFuture<AppSearchBatchResult<K, V>> future) {
+            mFuture = future;
+        }
+
+        @Override
+        public void onResult(@NonNull AppSearchBatchResult<K, V> result) {
+            mFuture.complete(result);
+        }
+
+        @Override
+        public void onSystemError(Throwable t) {
+            mFuture.completeExceptionally(t);
+        }
+    }
+}
diff --git a/services/tests/appfunctions/src/com/android/server/appfunctions/FutureAppSearchSessionTest.kt b/services/tests/appfunctions/src/com/android/server/appfunctions/FutureAppSearchSessionTest.kt
index 3bc44111ba08..edcbb9e81d87 100644
--- a/services/tests/appfunctions/src/com/android/server/appfunctions/FutureAppSearchSessionTest.kt
+++ b/services/tests/appfunctions/src/com/android/server/appfunctions/FutureAppSearchSessionTest.kt
@@ -21,6 +21,8 @@ import android.app.appfunctions.AppFunctionRuntimeMetadata.createAppFunctionRunt
 import android.app.appfunctions.AppFunctionRuntimeMetadata.createParentAppFunctionRuntimeSchema
 import android.app.appsearch.AppSearchBatchResult
 import android.app.appsearch.AppSearchManager
+import android.app.appsearch.GenericDocument
+import android.app.appsearch.GetByDocumentIdRequest
 import android.app.appsearch.PutDocumentsRequest
 import android.app.appsearch.RemoveByDocumentIdRequest
 import android.app.appsearch.SearchSpec
@@ -44,16 +46,16 @@ class FutureAppSearchSessionTest {
     @After
     fun clearData() {
         val searchContext = AppSearchManager.SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use {
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use {
             val setSchemaRequest = SetSchemaRequest.Builder().setForceOverride(true).build()
-            it.setSchema(setSchemaRequest)
+            it.setSchema(setSchemaRequest).get()
         }
     }
 
     @Test
     fun setSchema() {
         val searchContext = AppSearchManager.SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use { session ->
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use { session ->
             val setSchemaRequest =
                 SetSchemaRequest.Builder()
                     .addSchemas(
@@ -71,7 +73,7 @@ class FutureAppSearchSessionTest {
     @Test
     fun put() {
         val searchContext = AppSearchManager.SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use { session ->
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use { session ->
             val setSchemaRequest =
                 SetSchemaRequest.Builder()
                     .addSchemas(
@@ -97,7 +99,7 @@ class FutureAppSearchSessionTest {
     @Test
     fun remove() {
         val searchContext = AppSearchManager.SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use { session ->
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use { session ->
             val setSchemaRequest =
                 SetSchemaRequest.Builder()
                     .addSchemas(
@@ -131,7 +133,7 @@ class FutureAppSearchSessionTest {
     @Test
     fun search() {
         val searchContext = AppSearchManager.SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use { session ->
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use { session ->
             val setSchemaRequest =
                 SetSchemaRequest.Builder()
                     .addSchemas(
@@ -163,7 +165,7 @@ class FutureAppSearchSessionTest {
     @Test
     fun getByDocumentId() {
         val searchContext = AppSearchManager.SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use { session ->
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use { session ->
             val setSchemaRequest =
                 SetSchemaRequest.Builder()
                     .addSchemas(
@@ -171,24 +173,24 @@ class FutureAppSearchSessionTest {
                         createAppFunctionRuntimeSchema(TEST_PACKAGE_NAME),
                     )
                     .build()
-            val schema = session.setSchema(setSchemaRequest)
+            session.setSchema(setSchemaRequest).get()
             val appFunctionRuntimeMetadata =
                 AppFunctionRuntimeMetadata.Builder(TEST_PACKAGE_NAME, TEST_FUNCTION_ID, "").build()
             val putDocumentsRequest: PutDocumentsRequest =
                 PutDocumentsRequest.Builder()
                     .addGenericDocuments(appFunctionRuntimeMetadata)
                     .build()
-            val putResult = session.put(putDocumentsRequest)
+            session.put(putDocumentsRequest)
+            val getRequest =
+                GetByDocumentIdRequest.Builder(APP_FUNCTION_RUNTIME_NAMESPACE)
+                    .addIds(appFunctionRuntimeMetadata.id)
+                    .build()
 
-            val genricDocument =
-                session
-                    .getByDocumentId(
-                        /* documentId= */ "${TEST_PACKAGE_NAME}/${TEST_FUNCTION_ID}",
-                        APP_FUNCTION_RUNTIME_NAMESPACE,
-                    )
-                    .get()
+            val genericDocument: GenericDocument? =
+                session.getByDocumentId(getRequest).get().successes[appFunctionRuntimeMetadata.id]
 
-            val foundAppFunctionRuntimeMetadata = AppFunctionRuntimeMetadata(genricDocument)
+            assertThat(genericDocument).isNotNull()
+            val foundAppFunctionRuntimeMetadata = AppFunctionRuntimeMetadata(genericDocument!!)
             assertThat(foundAppFunctionRuntimeMetadata.functionId).isEqualTo(TEST_FUNCTION_ID)
         }
     }
diff --git a/services/tests/appfunctions/src/com/android/server/appfunctions/FutureGlobalSearchSessionTest.kt b/services/tests/appfunctions/src/com/android/server/appfunctions/FutureGlobalSearchSessionTest.kt
index 1fa55c7090aa..38cba6537c95 100644
--- a/services/tests/appfunctions/src/com/android/server/appfunctions/FutureGlobalSearchSessionTest.kt
+++ b/services/tests/appfunctions/src/com/android/server/appfunctions/FutureGlobalSearchSessionTest.kt
@@ -46,9 +46,9 @@ class FutureGlobalSearchSessionTest {
     @After
     fun clearData() {
         val searchContext = SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use {
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use {
             val setSchemaRequest = SetSchemaRequest.Builder().setForceOverride(true).build()
-            it.setSchema(setSchemaRequest)
+            it.setSchema(setSchemaRequest).get()
         }
     }
 
@@ -83,7 +83,7 @@ class FutureGlobalSearchSessionTest {
         assertThat(registerPackageObserver).isNull()
         // Trigger document change
         val searchContext = SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use { session ->
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use { session ->
             val setSchemaRequest =
                 SetSchemaRequest.Builder()
                     .addSchemas(
diff --git a/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt b/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
index 5769e07ee58d..3ebf68937674 100644
--- a/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
+++ b/services/tests/appfunctions/src/com/android/server/appfunctions/MetadataSyncAdapterTest.kt
@@ -41,9 +41,9 @@ class MetadataSyncAdapterTest {
     @After
     fun clearData() {
         val searchContext = SearchContext.Builder(TEST_DB).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use {
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use {
             val setSchemaRequest = SetSchemaRequest.Builder().setForceOverride(true).build()
-            it.setSchema(setSchemaRequest)
+            it.setSchema(setSchemaRequest).get()
         }
     }
 
@@ -61,7 +61,7 @@ class MetadataSyncAdapterTest {
                 .build()
         val putDocumentsRequest: PutDocumentsRequest =
             PutDocumentsRequest.Builder().addGenericDocuments(functionRuntimeMetadata).build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use {
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use {
             val setSchemaResponse = it.setSchema(setSchemaRequest).get()
             assertThat(setSchemaResponse).isNotNull()
             val appSearchBatchResult = it.put(putDocumentsRequest).get()
@@ -71,7 +71,7 @@ class MetadataSyncAdapterTest {
         val metadataSyncAdapter =
             MetadataSyncAdapter(
                 testExecutor,
-                FutureAppSearchSession(appSearchManager, testExecutor, searchContext),
+                FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext),
             )
         val packageToFunctionIdMap =
             metadataSyncAdapter.getPackageToFunctionIdMap(
@@ -111,7 +111,7 @@ class MetadataSyncAdapterTest {
                     functionRuntimeMetadata3,
                 )
                 .build()
-        FutureAppSearchSession(appSearchManager, testExecutor, searchContext).use {
+        FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext).use {
             val setSchemaResponse = it.setSchema(setSchemaRequest).get()
             assertThat(setSchemaResponse).isNotNull()
             val appSearchBatchResult = it.put(putDocumentsRequest).get()
@@ -121,7 +121,7 @@ class MetadataSyncAdapterTest {
         val metadataSyncAdapter =
             MetadataSyncAdapter(
                 testExecutor,
-                FutureAppSearchSession(appSearchManager, testExecutor, searchContext),
+                FutureAppSearchSessionImpl(appSearchManager, testExecutor, searchContext),
             )
         val packageToFunctionIdMap =
             metadataSyncAdapter.getPackageToFunctionIdMap(