Merge "Pass the root cause exception to the thrown exception on AppSearch failure" into udc-mainline-prod

4 files changed, 447 insertions, 36 deletions

diff --git a/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchConsentWorkerTest.java b/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchConsentWorkerTest.java
index e7344abe15..2bf1221c7a 100644
--- a/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchConsentWorkerTest.java
+++ b/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchConsentWorkerTest.java
@@ -64,6 +64,8 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.util.concurrent.FluentFuture;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
 
 import org.junit.Before;
 import org.junit.Test;
@@ -73,6 +75,9 @@ import org.mockito.Mockito;
 import java.util.Arrays;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 
 @SpyStatic(FlagsFactory.class)
 @RequiresSdkLevelAtLeastS
@@ -87,6 +92,8 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     private static final Topic TOPIC1 = Topic.create(0, 1, 11);
     private static final Topic TOPIC2 = Topic.create(12, 2, 22);
     private static final Topic TOPIC3 = Topic.create(123, 3, 33);
+    private static final int APPSEARCH_WRITE_TIMEOUT_MS = 1000;
+
     private final List<Topic> mTopics = Arrays.asList(TOPIC1, TOPIC2, TOPIC3);
 
     @Mock private Flags mMockFlags;
@@ -97,6 +104,8 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
         when(mMockFlags.getAdservicesApkShaCertificate())
                 .thenReturn(Flags.ADSERVICES_APK_SHA_CERTIFICATE);
         when(mMockFlags.getAppsearchWriterAllowListOverride()).thenReturn("");
+        // Reduce AppSearch write timeout to speed up the tests.
+        when(mMockFlags.getAppSearchWriteTimeout()).thenReturn(APPSEARCH_WRITE_TIMEOUT_MS);
     }
 
     @Test
@@ -144,6 +153,21 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
 
     @Test
     @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void testSetConsent_failure_timeout() {
+        initTimeoutResponse();
+
+        RuntimeException e =
+                assertThrows(
+                        RuntimeException.class,
+                        () -> AppSearchConsentWorker.getInstance().setConsent(API_TYPE, CONSENTED));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
+    @SpyStatic(PlatformStorage.class)
     @MockStatic(UserHandle.class)
     public void testSetConsent() {
         initSuccessResponse();
@@ -285,6 +309,8 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
                         RuntimeException.class,
                         () -> appSearchConsentWorker.clearAppsWithConsent(TEST));
         assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(ExecutionException.class);
     }
 
     @Test
@@ -408,6 +434,8 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
                         RuntimeException.class,
                         () -> appSearchConsentWorker.removeAppWithConsent(consentType, TEST));
         assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(ExecutionException.class);
     }
 
     @Test
@@ -488,7 +516,7 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
 
     @Test
     @SpyStatic(PlatformStorage.class)
-    public void testRecordGaUxNotificationDisplayed_faiure() {
+    public void testRecordGaUxNotificationDisplayed_failure() {
         runRecordNotificationDisplayedTestFailure(/* isBetaUx= */ false);
     }
 
@@ -511,6 +539,41 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     }
 
     @Test
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void testRecordNotificationDisplayed_failure_timeout() {
+        runRecordNotificationDisplayedTestFailureTimeout(/* isBetaUx= */ true);
+    }
+
+    @Test
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void testRecordGaUxNotificationDisplayed_failure_timeout() {
+        runRecordNotificationDisplayedTestFailureTimeout(/* isBetaUx= */ false);
+    }
+
+    private void runRecordNotificationDisplayedTestFailureTimeout(boolean isBetaUx) {
+        initTimeoutResponse();
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+
+        RuntimeException e;
+        if (isBetaUx) {
+            e =
+                    assertThrows(
+                            RuntimeException.class, () -> worker.recordNotificationDisplayed(true));
+        } else {
+            e =
+                    assertThrows(
+                            RuntimeException.class,
+                            () -> worker.recordGaUxNotificationDisplayed(true));
+        }
+
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
     @MockStatic(AppSearchNotificationDao.class)
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
@@ -574,6 +637,24 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @Test
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
+    public void testSetCurrentPrivacySandboxFeature_failure_timeout() {
+        initTimeoutResponse();
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        RuntimeException e =
+                assertThrows(
+                        RuntimeException.class,
+                        () ->
+                                worker.setCurrentPrivacySandboxFeature(
+                                        PrivacySandboxFeatureType.PRIVACY_SANDBOX_RECONSENT));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
     public void testSetCurrentPrivacySandboxFeature() {
         initSuccessResponse();
 
@@ -616,6 +697,27 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @MockStatic(AppSearchInteractionsDao.class)
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
+    public void testRecordUserManualInteractionWithConsent_failure_timeout() {
+        initTimeoutResponse();
+
+        when(AppSearchInteractionsDao.getRowId(any(), any())).thenReturn("" + UID);
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        int interactions = ConsentManager.MANUAL_INTERACTIONS_RECORDED;
+
+        RuntimeException e =
+                assertThrows(
+                        RuntimeException.class,
+                        () -> worker.recordUserManualInteractionWithConsent(interactions));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
+    @MockStatic(AppSearchInteractionsDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
     public void testRecordUserManualInteractionWithConsent() {
         initSuccessResponse();
 
@@ -638,7 +740,7 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
 
     @Test
     @SpyStatic(PlatformStorage.class)
-    public void testRecordBlockedTopics_failure() {
+    public void testRecordBlockedTopic_failure() {
         initFailureResponse();
 
         AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
@@ -651,7 +753,27 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @MockStatic(AppSearchTopicsConsentDao.class)
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
-    public void testRecordBlockedTopics_new() {
+    public void testRecordBlockedTopic_failure_timeout() {
+        initTimeoutResponse();
+
+        String query = "" + UID;
+        ExtendedMockito.doReturn(query).when(() -> AppSearchTopicsConsentDao.getQuery(any()));
+        ExtendedMockito.doReturn(null)
+                .when(() -> AppSearchTopicsConsentDao.readConsentData(any(), any(), any(), any()));
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        RuntimeException e =
+                assertThrows(RuntimeException.class, () -> worker.recordBlockedTopic(TOPIC1));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
+    @MockStatic(AppSearchTopicsConsentDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void testRecordBlockedTopic_new() {
         initSuccessResponse();
 
         String query = "" + UID;
@@ -667,7 +789,7 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @MockStatic(AppSearchTopicsConsentDao.class)
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
-    public void testRecordBlockedTopics() throws Exception {
+    public void testRecordBlockedTopic() {
         String query = "" + UID;
         ExtendedMockito.doReturn(query).when(() -> AppSearchTopicsConsentDao.getQuery(any()));
         AppSearchTopicsConsentDao dao = Mockito.mock(AppSearchTopicsConsentDao.class);
@@ -686,7 +808,7 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @Test
     @MockStatic(AppSearchTopicsConsentDao.class)
     @SpyStatic(PlatformStorage.class)
-    public void testRecordUnblockedTopics_failure() throws Exception {
+    public void testRecordUnblockedTopic_failure() throws Exception {
         AppSearchTopicsConsentDao dao = Mockito.mock(AppSearchTopicsConsentDao.class);
         ExtendedMockito.doReturn(dao)
                 .when(() -> AppSearchTopicsConsentDao.readConsentData(any(), any(), any(), any()));
@@ -702,13 +824,21 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
         verify(dao).removeBlockedTopic(TOPIC1);
         verify(dao).writeData(any(), any(), any());
         assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+
+        Throwable cause = e.getCause();
+        assertThat(cause).isNotNull();
+        assertThat(cause).isInstanceOf(ExecutionException.class);
+
+        Throwable rootCause = cause.getCause();
+        assertThat(rootCause).isNotNull();
+        assertThat(rootCause).isInstanceOf(InterruptedException.class);
     }
 
     @Test
     @MockStatic(AppSearchTopicsConsentDao.class)
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
-    public void testRecordUnblockedTopics_new() {
+    public void testRecordUnblockedTopic_new() {
         String query = "" + UID;
         ExtendedMockito.doReturn(query).when(() -> AppSearchTopicsConsentDao.getQuery(any()));
         ExtendedMockito.doReturn(null)
@@ -722,7 +852,7 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @MockStatic(AppSearchTopicsConsentDao.class)
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
-    public void testRecordUnblockedTopics() {
+    public void testRecordUnblockedTopic() {
         String query = "" + UID;
         ExtendedMockito.doReturn(query).when(() -> AppSearchTopicsConsentDao.getQuery(any()));
         AppSearchTopicsConsentDao dao = Mockito.mock(AppSearchTopicsConsentDao.class);
@@ -753,6 +883,20 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @Test
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
+    public void testClearBlockedTopics_failure_timeout() {
+        initTimeoutResponse();
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        RuntimeException e =
+                assertThrows(RuntimeException.class, () -> worker.clearBlockedTopics());
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
     public void testClearBlockedTopics() {
         initSuccessResponse();
 
@@ -781,6 +925,26 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
         when(mockResponse.getMigrationFailures()).thenReturn(List.of());
     }
 
+    private void initTimeoutResponse() {
+        AppSearchSession mockSession = Mockito.mock(AppSearchSession.class);
+        UserHandle mockUserHandle = Mockito.mock(UserHandle.class);
+        Mockito.when(UserHandle.getUserHandleForUid(Binder.getCallingUid()))
+                .thenReturn(mockUserHandle);
+        Mockito.when(mockUserHandle.getIdentifier()).thenReturn(UID);
+        ExtendedMockito.doReturn(Futures.immediateFuture(mockSession))
+                .when(() -> PlatformStorage.createSearchSessionAsync(any()));
+        verify(mockSession, atMost(1)).setSchemaAsync(any(SetSchemaRequest.class));
+
+        SetSchemaResponse mockResponse = Mockito.mock(SetSchemaResponse.class);
+        when(mockSession.setSchemaAsync(any(SetSchemaRequest.class)))
+                .thenReturn(Futures.immediateFuture(mockResponse));
+        AppSearchBatchResult<String, Void> result = Mockito.mock(AppSearchBatchResult.class);
+        when(mockSession.putAsync(any())).thenReturn(getLongRunningOperation(result));
+
+        verify(mockResponse, atMost(1)).getMigrationFailures();
+        when(mockResponse.getMigrationFailures()).thenReturn(List.of());
+    }
+
     private void initFailureResponse() {
         AppSearchSession mockSession = Mockito.mock(AppSearchSession.class);
         ExtendedMockito.doReturn(Futures.immediateFuture(mockSession))
@@ -801,6 +965,17 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
         when(mockResponse.getMigrationFailures()).thenReturn(List.of(failure));
     }
 
+    private <T> ListenableFuture<T> getLongRunningOperation(T result) {
+        // Wait for a time that's longer than the AppSearch write timeout, then return the result.
+        ListeningExecutorService ls =
+                MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor());
+        return ls.submit(
+                () -> {
+                    TimeUnit.MILLISECONDS.sleep(APPSEARCH_WRITE_TIMEOUT_MS + 500);
+                    return result;
+                });
+    }
+
     @Test
     @SpyStatic(AppSearchUxStatesDao.class)
     public void isAdIdEnabledTest_trueBit() {
@@ -826,8 +1001,6 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
     public void setAdIdEnabledTest_success() {
-        initSuccessResponse();
-
         String query = "" + UID;
         ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
         AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
@@ -845,17 +1018,17 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
 
     @Test
     @SpyStatic(PlatformStorage.class)
-    public void setAdIdEnabledTest_trueBit() {
-        setAdIdEnabledTest(true);
+    public void setAdIdEnabledTest_failure_trueBit() {
+        setAdIdEnabledTestFailure(true);
     }
 
     @Test
     @SpyStatic(PlatformStorage.class)
-    public void setAdIdEnabledTest_falseBit() {
-        setAdIdEnabledTest(false);
+    public void setAdIdEnabledTest_failure_falseBit() {
+        setAdIdEnabledTestFailure(false);
     }
 
-    private void setAdIdEnabledTest(boolean isAdIdEnabled) {
+    private void setAdIdEnabledTestFailure(boolean isAdIdEnabled) {
         initFailureResponse();
 
         AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
@@ -865,6 +1038,27 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     }
 
     @Test
+    @MockStatic(AppSearchUxStatesDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void setAdIdEnabledTest_timeout() {
+        String query = "" + UID;
+        ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
+        AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
+        ExtendedMockito.doReturn(dao)
+                .when(() -> AppSearchUxStatesDao.readData(any(), any(), any(), any()));
+        when(dao.writeData(any(), any(), any()))
+                .thenReturn(FluentFuture.from(getLongRunningOperation(null)));
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        RuntimeException e =
+                assertThrows(RuntimeException.class, () -> worker.setAdIdEnabled(true));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
     @SpyStatic(AppSearchUxStatesDao.class)
     public void isU18AccountTest_trueBit() {
         isU18AccountTest(true);
@@ -889,8 +1083,6 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
     public void setU18AccountTest_success() {
-        initSuccessResponse();
-
         String query = "" + UID;
         ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
         AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
@@ -928,6 +1120,28 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     }
 
     @Test
+    @MockStatic(AppSearchUxStatesDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void setU18AccountTest_timeout() {
+        String query = "" + UID;
+        ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
+        AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
+        ExtendedMockito.doReturn(dao)
+                .when(() -> AppSearchUxStatesDao.readData(any(), any(), any(), any()));
+        AppSearchBatchResult<String, Void> result = Mockito.mock(AppSearchBatchResult.class);
+        when(dao.writeData(any(), any(), any()))
+                .thenReturn(FluentFuture.from(getLongRunningOperation(result)));
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        RuntimeException e =
+                assertThrows(RuntimeException.class, () -> worker.setU18Account(false));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
     @SpyStatic(AppSearchUxStatesDao.class)
     public void isEntryPointEnabledTest_trueBit() {
         isEntryPointEnabledTest(true);
@@ -996,6 +1210,30 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     }
 
     @Test
+    @MockStatic(AppSearchUxStatesDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void setEntryPointEnabledTest_timeout() {
+        initSuccessResponse();
+
+        String query = "" + UID;
+        ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
+        AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
+        ExtendedMockito.doReturn(dao)
+                .when(() -> AppSearchUxStatesDao.readData(any(), any(), any(), any()));
+        AppSearchBatchResult<String, Void> result = Mockito.mock(AppSearchBatchResult.class);
+        when(dao.writeData(any(), any(), any()))
+                .thenReturn(FluentFuture.from(getLongRunningOperation(result)));
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        RuntimeException e =
+                assertThrows(RuntimeException.class, () -> worker.setEntryPointEnabled(true));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
     @SpyStatic(AppSearchUxStatesDao.class)
     public void isAdultAccountTest_trueBit() {
         isAdultAccountTest(true);
@@ -1040,7 +1278,7 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @MockStatic(AppSearchUxStatesDao.class)
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
-    public void setU18AdultAccountTest_success() {
+    public void setAdultAccountTest_success() {
         initSuccessResponse();
 
         String query = "" + UID;
@@ -1059,6 +1297,30 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     }
 
     @Test
+    @MockStatic(AppSearchUxStatesDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void setAdultAccountTest_timeout() {
+        initSuccessResponse();
+
+        String query = "" + UID;
+        ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
+        AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
+        ExtendedMockito.doReturn(dao)
+                .when(() -> AppSearchUxStatesDao.readData(any(), any(), any(), any()));
+        AppSearchBatchResult<String, Void> result = Mockito.mock(AppSearchBatchResult.class);
+        when(dao.writeData(any(), any(), any()))
+                .thenReturn(FluentFuture.from(getLongRunningOperation(result)));
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        RuntimeException e =
+                assertThrows(RuntimeException.class, () -> worker.setAdultAccount(true));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
     @SpyStatic(AppSearchUxStatesDao.class)
     public void wasU18NotificationDisplayedTest_trueBit() {
         wasU18NotificationDisplayedTest(true);
@@ -1128,6 +1390,31 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     }
 
     @Test
+    @MockStatic(AppSearchUxStatesDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void setU18NotificationDisplayedTest_timeout() {
+        initSuccessResponse();
+
+        String query = "" + UID;
+        ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
+        AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
+        ExtendedMockito.doReturn(dao)
+                .when(() -> AppSearchUxStatesDao.readData(any(), any(), any(), any()));
+        AppSearchBatchResult<String, Void> result = Mockito.mock(AppSearchBatchResult.class);
+        when(dao.writeData(any(), any(), any()))
+                .thenReturn(FluentFuture.from(getLongRunningOperation(result)));
+
+        AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+        RuntimeException e =
+                assertThrows(
+                        RuntimeException.class, () -> worker.setU18NotificationDisplayed(true));
+        assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+        assertThat(e.getCause()).isNotNull();
+        assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+    }
+
+    @Test
     @SpyStatic(AppSearchUxStatesDao.class)
     public void getUxTest_allUxs() {
         for (PrivacySandboxUxCollection ux : PrivacySandboxUxCollection.values()) {
@@ -1155,8 +1442,6 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
     public void setUxTest_allUxsSuccess() {
-        initSuccessResponse();
-
         for (PrivacySandboxUxCollection ux : PrivacySandboxUxCollection.values()) {
             String query = "" + UID;
             ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
@@ -1175,6 +1460,29 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     }
 
     @Test
+    @MockStatic(AppSearchUxStatesDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void setUxTest_allUxs_timeout() {
+        for (PrivacySandboxUxCollection ux : PrivacySandboxUxCollection.values()) {
+            String query = "" + UID;
+            ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
+            AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
+            ExtendedMockito.doReturn(dao)
+                    .when(() -> AppSearchUxStatesDao.readData(any(), any(), any(), any()));
+            AppSearchBatchResult<String, Void> result = Mockito.mock(AppSearchBatchResult.class);
+            when(dao.writeData(any(), any(), any()))
+                    .thenReturn(FluentFuture.from(getLongRunningOperation(result)));
+
+            AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+            RuntimeException e = assertThrows(RuntimeException.class, () -> worker.setUx(ux));
+            assertThat(e.getMessage()).isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+            assertThat(e.getCause()).isNotNull();
+            assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+        }
+    }
+
+    @Test
     @SpyStatic(AppSearchUxStatesDao.class)
     public void getEnrollmentChannelTest_allUxsAllEnrollmentChannels() {
         for (PrivacySandboxUxCollection ux : PrivacySandboxUxCollection.values()) {
@@ -1217,8 +1525,6 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
     @SpyStatic(PlatformStorage.class)
     @SpyStatic(UserHandle.class)
     public void setEnrollmentChannelTest_allUxsAllEnrollmentChannelsSuccess() {
-        initSuccessResponse();
-
         for (PrivacySandboxUxCollection ux : PrivacySandboxUxCollection.values()) {
             for (PrivacySandboxEnrollmentChannelCollection channel :
                     ux.getEnrollmentChannelCollection()) {
@@ -1240,4 +1546,35 @@ public final class AppSearchConsentWorkerTest extends AdServicesExtendedMockitoT
             }
         }
     }
+
+    @Test
+    @MockStatic(AppSearchUxStatesDao.class)
+    @SpyStatic(PlatformStorage.class)
+    @SpyStatic(UserHandle.class)
+    public void setEnrollmentChannelTest_allUxsAllEnrollmentChannels_timeout() {
+        for (PrivacySandboxUxCollection ux : PrivacySandboxUxCollection.values()) {
+            for (PrivacySandboxEnrollmentChannelCollection channel :
+                    ux.getEnrollmentChannelCollection()) {
+                String query = "" + UID;
+                ExtendedMockito.doReturn(query).when(() -> AppSearchUxStatesDao.getQuery(any()));
+                AppSearchUxStatesDao dao = Mockito.mock(AppSearchUxStatesDao.class);
+                ExtendedMockito.doReturn(dao)
+                        .when(() -> AppSearchUxStatesDao.readData(any(), any(), any(), any()));
+                AppSearchBatchResult<String, Void> result =
+                        Mockito.mock(AppSearchBatchResult.class);
+                when(dao.writeData(any(), any(), any()))
+                        .thenReturn(FluentFuture.from(getLongRunningOperation(result)));
+
+                AppSearchConsentWorker worker = AppSearchConsentWorker.getInstance();
+                RuntimeException e =
+                        assertThrows(
+                                RuntimeException.class,
+                                () -> worker.setEnrollmentChannel(ux, channel));
+                assertThat(e.getMessage())
+                        .isEqualTo(ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+                assertThat(e.getCause()).isNotNull();
+                assertThat(e.getCause()).isInstanceOf(TimeoutException.class);
+            }
+        }
+    }
 }
diff --git a/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchDaoTest.java b/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchDaoTest.java
index 47bc186455..116891316c 100644
--- a/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchDaoTest.java
+++ b/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchDaoTest.java
@@ -52,6 +52,8 @@ import com.android.adservices.service.consent.ConsentConstants;
 import com.google.common.util.concurrent.FluentFuture;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
+import com.google.common.util.concurrent.ListeningExecutorService;
+import com.google.common.util.concurrent.MoreExecutors;
 
 import org.junit.Before;
 import org.junit.Rule;
@@ -63,6 +65,8 @@ import org.mockito.MockitoAnnotations;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 @SmallTest
 public class AppSearchDaoTest {
@@ -87,6 +91,8 @@ public class AppSearchDaoTest {
             new PackageIdentifier(
                     /* packageName= */ TEST, /* sha256= */ new Signature(SHA).toByteArray());
 
+    private static final int APPSEARCH_READ_TIMEOUT_MS = 500;
+
     @Rule
     public final AdServicesExtendedMockitoRule adServicesExtendedMockitoRule =
             new AdServicesExtendedMockitoRule.Builder(this).mockStatic(FlagsFactory.class).build();
@@ -95,6 +101,7 @@ public class AppSearchDaoTest {
     public void before() {
         MockitoAnnotations.initMocks(this);
         when(mFlags.getAppsearchWriterAllowListOverride()).thenReturn("");
+        when(mFlags.getAppSearchReadTimeout()).thenReturn(APPSEARCH_READ_TIMEOUT_MS);
         doReturn(mFlags).when(FlagsFactory::getFlags);
     }
 
@@ -205,6 +212,30 @@ public class AppSearchDaoTest {
     }
 
     @Test
+    public void testReadConsentData_timeout() {
+        AppSearchDao result =
+                AppSearchDao.readConsentData(
+                        AppSearchConsentDao.class,
+                        getLongRunningOperation(mGlobalSearchSession),
+                        mExecutor,
+                        NAMESPACE,
+                        TEST,
+                        mAdServicesPackageName);
+        assertThat(result).isNull();
+    }
+
+    private <T> ListenableFuture<T> getLongRunningOperation(T result) {
+        // Wait for a time that's longer than the AppSearch read timeout, then return the result.
+        ListeningExecutorService ls =
+                MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor());
+        return ls.submit(
+                () -> {
+                    TimeUnit.MILLISECONDS.sleep(APPSEARCH_READ_TIMEOUT_MS + 500);
+                    return result;
+                });
+    }
+
+    @Test
     public void testReadAppSearchData_emptyQuery() {
         AppSearchDao dao =
                 AppSearchDao.readAppSearchSessionData(
@@ -262,7 +293,8 @@ public class AppSearchDaoTest {
         when(mockSession.setSchemaAsync(any(SetSchemaRequest.class)))
                 .thenReturn(Futures.immediateFuture(mockResponse));
 
-        AppSearchResult mockResult = Mockito.mock(AppSearchResult.class);
+        AppSearchResult<Void> mockResult =
+                AppSearchResult.newFailedResult(AppSearchResult.RESULT_INVALID_ARGUMENT, "test");
         SetSchemaResponse.MigrationFailure failure =
                 new SetSchemaResponse.MigrationFailure(
                         /* namespace= */ TEST,
@@ -278,11 +310,12 @@ public class AppSearchDaoTest {
                         Futures.immediateFuture(mockSession),
                         List.of(PACKAGE_IDENTIFIER),
                         mExecutor);
-        ExecutionException e = assertThrows(ExecutionException.class, () -> result.get());
+        ExecutionException e = assertThrows(ExecutionException.class, result::get);
         assertThat(e.getMessage())
                 .isEqualTo(
                         "java.lang.RuntimeException: "
-                                + ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE);
+                                + ConsentConstants.ERROR_MESSAGE_APPSEARCH_FAILURE
+                                + " Migration failure: [FAILURE(3)]: test");
     }
 
     @Test
@@ -337,7 +370,7 @@ public class AppSearchDaoTest {
                         mExecutor,
                         TEST,
                         NAMESPACE);
-        ExecutionException e = assertThrows(ExecutionException.class, () -> result.get());
+        ExecutionException e = assertThrows(ExecutionException.class, result::get);
         assertThat(e.getMessage())
                 .isEqualTo(
                         "java.lang.RuntimeException: "
diff --git a/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchMeasurementRollbackWorkerTest.java b/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchMeasurementRollbackWorkerTest.java
index 0765ba5275..638ee67142 100644
--- a/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchMeasurementRollbackWorkerTest.java
+++ b/adservices/tests/unittest/service-core/appsearch/src/com/android/adservices/service/appsearch/AppSearchMeasurementRollbackWorkerTest.java
@@ -41,6 +41,8 @@ import androidx.test.filters.SmallTest;
 
 import com.android.adservices.concurrency.AdServicesExecutors;
 import com.android.adservices.mockito.AdServicesExtendedMockitoRule;
+import com.android.adservices.service.Flags;
+import com.android.adservices.service.FlagsFactory;
 import com.android.adservices.service.common.compat.FileCompatUtils;
 import com.android.adservices.service.consent.ConsentConstants;
 
@@ -66,7 +68,7 @@ public class AppSearchMeasurementRollbackWorkerTest {
             FileCompatUtils.getAdservicesFilename("measurement_rollback");
     private static final String USERID = "user1";
     private static final long APEX_VERSION = 100L;
-    private static final int FUTURE_TIMEOUT_MILLISECONDS = 3000;
+    private static final int APPSEARCH_WRITE_TIMEOUT_MS = 1000;
 
     private final Context mContext = ApplicationProvider.getApplicationContext();
     private final String mAdServicesPackageName =
@@ -74,17 +76,22 @@ public class AppSearchMeasurementRollbackWorkerTest {
     private final Executor mExecutor = AdServicesExecutors.getBackgroundExecutor();
     private AppSearchMeasurementRollbackWorker mWorker;
     @Mock private ListenableFuture<AppSearchSession> mAppSearchSession;
+    @Mock private Flags mMockFlags;
 
     @Rule
     public final AdServicesExtendedMockitoRule adServicesExtendedMockitoRule =
             new AdServicesExtendedMockitoRule.Builder(this)
                     .mockStatic(PlatformStorage.class)
                     .mockStatic(AppSearchDao.class)
+                    .mockStatic(FlagsFactory.class)
                     .setStrictness(Strictness.LENIENT)
                     .build();
 
     @Before
     public void setup() {
+        doReturn(mMockFlags).when(FlagsFactory::getFlags);
+        doReturn(APPSEARCH_WRITE_TIMEOUT_MS).when(mMockFlags).getAppSearchWriteTimeout();
+
         ArgumentCaptor<PlatformStorage.SearchContext> cap =
                 ArgumentCaptor.forClass(PlatformStorage.SearchContext.class);
         doReturn(mAppSearchSession)
@@ -107,7 +114,7 @@ public class AppSearchMeasurementRollbackWorkerTest {
     @SuppressWarnings("FutureReturnValueIgnored")
     @Test
     public void testClearAdServicesDeletionOccurred() {
-        FluentFuture mockResult =
+        FluentFuture<AppSearchBatchResult<String, Void>> mockResult =
                 FluentFuture.from(
                         Futures.immediateFuture(
                                 new AppSearchBatchResult.Builder<String, Void>().build()));
@@ -128,13 +135,14 @@ public class AppSearchMeasurementRollbackWorkerTest {
 
     @Test
     public void testClearAdServicesDeletionOccurred_throwsChecked() {
-        Callable<Void> callable =
+        Callable<AppSearchBatchResult<String, Void>> callable =
                 () -> {
-                    TimeUnit.MILLISECONDS.sleep(FUTURE_TIMEOUT_MILLISECONDS);
+                    TimeUnit.MILLISECONDS.sleep(APPSEARCH_WRITE_TIMEOUT_MS + 500);
                     return null;
                 };
 
-        FluentFuture mockResult = FluentFuture.from(Futures.submit(callable, mExecutor));
+        FluentFuture<AppSearchBatchResult<String, Void>> mockResult =
+                FluentFuture.from(Futures.submit(callable, mExecutor));
         doReturn(mockResult).when(() -> AppSearchDao.deleteData(any(), any(), any(), any(), any()));
 
         RuntimeException e =
@@ -234,17 +242,17 @@ public class AppSearchMeasurementRollbackWorkerTest {
 
     @Test
     public void testRecordAdServicesDeletionOccurred_throwsChecked() {
-        // The manager class waits for 2 seconds on the future.get() call before timing out. So
-        // creating a future that takes longer than 2 sec to resolve, in order to create a
+        // The manager class waits for a few seconds on the future.get() call before timing out. So
+        // creating a future that takes longer than the timeout to resolve, in order to create a
         // TimeoutException.
-        Callable<Void> callable =
+        Callable<AppSearchBatchResult<String, Void>> callable =
                 () -> {
-                    TimeUnit.MILLISECONDS.sleep(FUTURE_TIMEOUT_MILLISECONDS);
+                    TimeUnit.MILLISECONDS.sleep(APPSEARCH_WRITE_TIMEOUT_MS + 500);
                     return null;
                 };
 
-        ListenableFuture<Void> future = Futures.submit(callable, mExecutor);
-        FluentFuture mockFuture = FluentFuture.from(future);
+        FluentFuture<AppSearchBatchResult<String, Void>> mockFuture =
+                FluentFuture.from(Futures.submit(callable, mExecutor));
         AppSearchMeasurementRollbackDao dao = mock(AppSearchMeasurementRollbackDao.class);
         doReturn(mockFuture).when(dao).writeData(any(), any(), any());
 
diff --git a/adservices/tests/unittest/service-core/src/com/android/adservices/service/PhFlagsTest.java b/adservices/tests/unittest/service-core/src/com/android/adservices/service/PhFlagsTest.java
index 25aacd5191..f86269555f 100644
--- a/adservices/tests/unittest/service-core/src/com/android/adservices/service/PhFlagsTest.java
+++ b/adservices/tests/unittest/service-core/src/com/android/adservices/service/PhFlagsTest.java
@@ -50,6 +50,8 @@ import static com.android.adservices.service.Flags.DEFAULT_ADSERVICES_CONSENT_MI
 import static com.android.adservices.service.Flags.DEFAULT_ADSERVICES_ENABLEMENT_CHECK_ENABLED;
 import static com.android.adservices.service.Flags.DEFAULT_ADSERVICES_VERSION_MAPPINGS;
 import static com.android.adservices.service.Flags.DEFAULT_AD_ID_FETCHER_TIMEOUT_MS;
+import static com.android.adservices.service.Flags.DEFAULT_APPSEARCH_READ_TIMEOUT_MS;
+import static com.android.adservices.service.Flags.DEFAULT_APPSEARCH_WRITE_TIMEOUT_MS;
 import static com.android.adservices.service.Flags.DEFAULT_AUCTION_SERVER_AD_ID_FETCHER_TIMEOUT_MS;
 import static com.android.adservices.service.Flags.DEFAULT_BACKGROUND_JOB_SAMPLING_LOGGING_RATE;
 import static com.android.adservices.service.Flags.DEFAULT_BLOCKED_TOPICS_SOURCE_OF_TRUTH;
@@ -9965,6 +9967,37 @@ public class PhFlagsTest {
         assertThrows(IllegalArgumentException.class, mPhFlags::getBackgroundJobSamplingLoggingRate);
     }
 
+    @Test
+    public void testGetAppSearchWriteTimeout() {
+        // Without any overriding, the value is the hard coded constant.
+        assertThat(mPhFlags.getAppSearchWriteTimeout())
+                .isEqualTo(DEFAULT_APPSEARCH_WRITE_TIMEOUT_MS);
+
+        int phOverridingValue = DEFAULT_APPSEARCH_WRITE_TIMEOUT_MS + 1000;
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_ADSERVICES,
+                FlagsConstants.KEY_APPSEARCH_WRITE_TIMEOUT_MS,
+                Integer.toString(phOverridingValue),
+                /* makeDefault */ false);
+
+        assertThat(mPhFlags.getAppSearchWriteTimeout()).isEqualTo(phOverridingValue);
+    }
+
+    @Test
+    public void testGetAppSearchReadTimeout() {
+        // Without any overriding, the value is the hard coded constant.
+        assertThat(mPhFlags.getAppSearchReadTimeout()).isEqualTo(DEFAULT_APPSEARCH_READ_TIMEOUT_MS);
+
+        int phOverridingValue = DEFAULT_APPSEARCH_READ_TIMEOUT_MS + 500;
+        DeviceConfig.setProperty(
+                DeviceConfig.NAMESPACE_ADSERVICES,
+                FlagsConstants.KEY_APPSEARCH_READ_TIMEOUT_MS,
+                Integer.toString(phOverridingValue),
+                /* makeDefault */ false);
+
+        assertThat(mPhFlags.getAppSearchReadTimeout()).isEqualTo(phOverridingValue);
+    }
+
     private void overrideBackgroundJobSamplingLoggingRate(int phOverridingValue) {
         DeviceConfig.setProperty(
                 DeviceConfig.NAMESPACE_ADSERVICES,