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

2 files changed, 229 insertions, 174 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);
+        }
+    }
+}