1 /*
2  * Copyright (C) 2022 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package com.android.adservices.service.customaudience;
18 
19 import static android.adservices.common.AdServicesStatusUtils.STATUS_INTERNAL_ERROR;
20 import static android.adservices.common.AdServicesStatusUtils.STATUS_SUCCESS;
21 
22 import static com.android.adservices.common.CommonFlagsValues.EXTENDED_FLEDGE_BACKGROUND_FETCH_NETWORK_CONNECT_TIMEOUT_MS;
23 import static com.android.adservices.common.CommonFlagsValues.EXTENDED_FLEDGE_BACKGROUND_FETCH_NETWORK_READ_TIMEOUT_MS;
24 
25 import static com.google.common.truth.Truth.assertThat;
26 import static com.google.common.util.concurrent.Futures.immediateFuture;
27 import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
28 
29 import static org.junit.Assert.assertThrows;
30 import static org.junit.Assert.assertTrue;
31 import static org.mockito.ArgumentMatchers.anyInt;
32 import static org.mockito.Mockito.any;
33 import static org.mockito.Mockito.anyLong;
34 import static org.mockito.Mockito.doAnswer;
35 import static org.mockito.Mockito.doReturn;
36 import static org.mockito.Mockito.doThrow;
37 import static org.mockito.Mockito.never;
38 import static org.mockito.Mockito.times;
39 import static org.mockito.Mockito.verify;
40 import static org.mockito.Mockito.when;
41 
42 import android.adservices.common.CommonFixture;
43 import android.annotation.NonNull;
44 import android.content.Context;
45 import android.content.pm.PackageManager;
46 
47 import androidx.room.Room;
48 import androidx.test.core.app.ApplicationProvider;
49 import androidx.test.filters.FlakyTest;
50 
51 import com.android.adservices.LoggerFactory;
52 import com.android.adservices.common.AdServicesDeviceSupportedRule;
53 import com.android.adservices.concurrency.AdServicesExecutors;
54 import com.android.adservices.customaudience.DBCustomAudienceBackgroundFetchDataFixture;
55 import com.android.adservices.data.adselection.AppInstallDao;
56 import com.android.adservices.data.adselection.SharedStorageDatabase;
57 import com.android.adservices.data.customaudience.CustomAudienceDao;
58 import com.android.adservices.data.customaudience.CustomAudienceDatabase;
59 import com.android.adservices.data.customaudience.DBCustomAudience;
60 import com.android.adservices.data.customaudience.DBCustomAudienceBackgroundFetchData;
61 import com.android.adservices.data.enrollment.EnrollmentDao;
62 import com.android.adservices.service.FakeFlagsFactory;
63 import com.android.adservices.service.Flags;
64 import com.android.adservices.service.customaudience.BackgroundFetchRunner.UpdateResultType;
65 import com.android.adservices.service.stats.AdServicesLoggerImpl;
66 import com.android.adservices.service.stats.BackgroundFetchExecutionLogger;
67 import com.android.adservices.service.stats.CustomAudienceLoggerFactory;
68 import com.android.adservices.service.stats.UpdateCustomAudienceExecutionLogger;
69 import com.android.adservices.shared.testing.AnswerSyncCallback;
70 import com.android.adservices.shared.testing.SdkLevelSupportRule;
71 import com.android.adservices.shared.testing.concurrency.SimpleSyncCallback;
72 
73 import com.google.common.util.concurrent.FluentFuture;
74 import com.google.common.util.concurrent.ListenableFuture;
75 
76 import org.junit.Before;
77 import org.junit.Rule;
78 import org.junit.Test;
79 import org.mockito.Mock;
80 import org.mockito.Mockito;
81 import org.mockito.junit.MockitoJUnit;
82 import org.mockito.junit.MockitoRule;
83 
84 import java.time.Clock;
85 import java.time.Instant;
86 import java.util.ArrayList;
87 import java.util.Arrays;
88 import java.util.List;
89 import java.util.concurrent.CountDownLatch;
90 import java.util.concurrent.ExecutionException;
91 import java.util.concurrent.ExecutorService;
92 import java.util.concurrent.Executors;
93 import java.util.concurrent.TimeUnit;
94 import java.util.concurrent.TimeoutException;
95 import java.util.concurrent.atomic.AtomicInteger;
96 
97 public class BackgroundFetchWorkerTest {
98     private static final LoggerFactory.Logger sLogger = LoggerFactory.getFledgeLogger();
99     private static final Context CONTEXT = ApplicationProvider.getApplicationContext();
100 
101     private final Flags mFlags = new BackgroundFetchWorkerTestFlags(true);
102     private final ExecutorService mExecutorService = Executors.newFixedThreadPool(8);
103 
104     @Rule(order = 0)
105     public final SdkLevelSupportRule sdkLevel = SdkLevelSupportRule.forAtLeastS();
106 
107     @Rule(order = 1)
108     public final AdServicesDeviceSupportedRule deviceSupportRule =
109             new AdServicesDeviceSupportedRule();
110 
111     @Rule(order = 2)
112     public final MockitoRule mockitoRule = MockitoJUnit.rule();
113 
114     @Mock private PackageManager mPackageManagerMock;
115     @Mock private EnrollmentDao mEnrollmentDaoMock;
116     @Mock private Clock mClockMock;
117     @Mock private AdServicesLoggerImpl mAdServicesLoggerImplMock;
118     @Mock private UpdateCustomAudienceExecutionLogger mUpdateCustomAudienceExecutionLoggerMock;
119     private BackgroundFetchExecutionLogger mBackgroundFetchExecutionLoggerSpy;
120     @Mock private CustomAudienceLoggerFactory mCustomAudienceLoggerFactoryMock;
121 
122     private CustomAudienceDao mCustomAudienceDaoSpy;
123     private AppInstallDao mAppInstallDaoSpy;
124     private BackgroundFetchRunner mBackgroundFetchRunnerSpy;
125     private BackgroundFetchWorker mBackgroundFetchWorker;
126 
127     @Before
setup()128     public void setup() {
129         when(mCustomAudienceLoggerFactoryMock.getUpdateCustomAudienceExecutionLogger())
130                 .thenReturn(mUpdateCustomAudienceExecutionLoggerMock);
131         mBackgroundFetchExecutionLoggerSpy =
132                 Mockito.spy(
133                         new BackgroundFetchExecutionLogger(
134                                 com.android.adservices.shared.util.Clock.getInstance(),
135                                 mAdServicesLoggerImplMock));
136         when(mCustomAudienceLoggerFactoryMock.getBackgroundFetchExecutionLogger())
137                 .thenReturn(mBackgroundFetchExecutionLoggerSpy);
138 
139         mCustomAudienceDaoSpy =
140                 Mockito.spy(
141                         Room.inMemoryDatabaseBuilder(CONTEXT, CustomAudienceDatabase.class)
142                                 .addTypeConverter(new DBCustomAudience.Converters(true, true, true))
143                                 .build()
144                                 .customAudienceDao());
145         mAppInstallDaoSpy =
146                 Mockito.spy(
147                         Room.inMemoryDatabaseBuilder(CONTEXT, SharedStorageDatabase.class)
148                                 .build()
149                                 .appInstallDao());
150 
151         mBackgroundFetchRunnerSpy =
152                 Mockito.spy(
153                         new BackgroundFetchRunner(
154                                 mCustomAudienceDaoSpy,
155                                 mAppInstallDaoSpy,
156                                 mPackageManagerMock,
157                                 mEnrollmentDaoMock,
158                                 mFlags,
159                                 mCustomAudienceLoggerFactoryMock));
160 
161         mBackgroundFetchWorker =
162                 new BackgroundFetchWorker(
163                         mCustomAudienceDaoSpy,
164                         mFlags,
165                         mBackgroundFetchRunnerSpy,
166                         mClockMock,
167                         mCustomAudienceLoggerFactoryMock);
168     }
169 
170     @Test
testBackgroundFetchWorkerNullInputsCauseFailure()171     public void testBackgroundFetchWorkerNullInputsCauseFailure() {
172         assertThrows(
173                 NullPointerException.class,
174                 () ->
175                         new BackgroundFetchWorker(
176                                 null,
177                                 FakeFlagsFactory.getFlagsForTest(),
178                                 mBackgroundFetchRunnerSpy,
179                                 mClockMock,
180                                 mCustomAudienceLoggerFactoryMock));
181 
182         assertThrows(
183                 NullPointerException.class,
184                 () ->
185                         new BackgroundFetchWorker(
186                                 mCustomAudienceDaoSpy,
187                                 null,
188                                 mBackgroundFetchRunnerSpy,
189                                 mClockMock,
190                                 mCustomAudienceLoggerFactoryMock));
191 
192         assertThrows(
193                 NullPointerException.class,
194                 () ->
195                         new BackgroundFetchWorker(
196                                 mCustomAudienceDaoSpy,
197                                 FakeFlagsFactory.getFlagsForTest(),
198                                 null,
199                                 mClockMock,
200                                 mCustomAudienceLoggerFactoryMock));
201 
202         assertThrows(
203                 NullPointerException.class,
204                 () ->
205                         new BackgroundFetchWorker(
206                                 mCustomAudienceDaoSpy,
207                                 FakeFlagsFactory.getFlagsForTest(),
208                                 mBackgroundFetchRunnerSpy,
209                                 null,
210                                 mCustomAudienceLoggerFactoryMock));
211         assertThrows(
212                 NullPointerException.class,
213                 () ->
214                         new BackgroundFetchWorker(
215                                 mCustomAudienceDaoSpy,
216                                 FakeFlagsFactory.getFlagsForTest(),
217                                 mBackgroundFetchRunnerSpy,
218                                 mClockMock,
219                                 null));
220     }
221 
222     @Test
testRunBackgroundFetchThrowsTimeoutDuringUpdates()223     public void testRunBackgroundFetchThrowsTimeoutDuringUpdates() throws InterruptedException {
224         class FlagsWithSmallTimeout implements Flags {
225             @Override
226             public long getFledgeBackgroundFetchJobMaxRuntimeMs() {
227                 return 100L;
228             }
229         }
230 
231         class BackgroundFetchRunnerWithSleep extends BackgroundFetchRunner {
232             BackgroundFetchRunnerWithSleep(
233                     @NonNull CustomAudienceDao customAudienceDao, @NonNull Flags flags) {
234                 super(
235                         customAudienceDao,
236                         mAppInstallDaoSpy,
237                         mPackageManagerMock,
238                         mEnrollmentDaoMock,
239                         flags,
240                         mCustomAudienceLoggerFactoryMock);
241             }
242 
243             @Override
244             public void deleteExpiredCustomAudiences(@NonNull Instant jobStartTime) {
245                 // Do nothing
246             }
247 
248             @Override
249             public FluentFuture<UpdateResultType> updateCustomAudience(
250                     @NonNull Instant jobStartTime,
251                     @NonNull DBCustomAudienceBackgroundFetchData fetchData) {
252 
253                 return FluentFuture.from(
254                         AdServicesExecutors.getBlockingExecutor()
255                                 .submit(
256                                         () -> {
257                                             try {
258                                                 Thread.sleep(500L);
259                                             } catch (InterruptedException e) {
260                                                 e.printStackTrace();
261                                             }
262                                             return null;
263                                         }));
264             }
265         }
266 
267         Flags flagsWithSmallTimeout = new FlagsWithSmallTimeout();
268         BackgroundFetchRunner backgroundFetchRunnerWithSleep =
269                 new BackgroundFetchRunnerWithSleep(mCustomAudienceDaoSpy, flagsWithSmallTimeout);
270         BackgroundFetchWorker backgroundFetchWorkerThatTimesOut =
271                 new BackgroundFetchWorker(
272                         mCustomAudienceDaoSpy,
273                         flagsWithSmallTimeout,
274                         backgroundFetchRunnerWithSleep,
275                         mClockMock,
276                         mCustomAudienceLoggerFactoryMock);
277 
278         // Mock a custom audience eligible for update
279         DBCustomAudienceBackgroundFetchData fetchData =
280                 DBCustomAudienceBackgroundFetchDataFixture.getValidBuilderByBuyer(
281                                 CommonFixture.VALID_BUYER_1)
282                         .setEligibleUpdateTime(CommonFixture.FIXED_NOW)
283                         .build();
284         List<DBCustomAudienceBackgroundFetchData> fetchDataList = Arrays.asList(fetchData);
285         doReturn(fetchDataList)
286                 .when(mCustomAudienceDaoSpy)
287                 .getActiveEligibleCustomAudienceBackgroundFetchData(any(), anyLong());
288 
289         when(mClockMock.instant()).thenReturn(Instant.now());
290 
291         // Ensure that log is closed after timeout.
292         CountDownLatch closeLoggerLatch = new CountDownLatch(1);
293         doAnswer(
294                         unusedInovcation -> {
295                             closeLoggerLatch.countDown();
296                             return null;
297                         })
298                 .when(mBackgroundFetchExecutionLoggerSpy)
299                 .close(anyInt(), anyInt());
300 
301         // Time out while updating custom audiences
302         ExecutionException expected =
303                 assertThrows(
304                         ExecutionException.class,
305                         () -> backgroundFetchWorkerThatTimesOut.runBackgroundFetch().get());
306         assertThat(expected.getCause()).isInstanceOf(TimeoutException.class);
307         // Wait for logger to close.
308         assertThat(closeLoggerLatch.await(200, TimeUnit.MILLISECONDS)).isTrue();
309         // Background data already fetched, internal error due to failure in runner.
310         verify(mBackgroundFetchExecutionLoggerSpy)
311                 .close(fetchDataList.size(), STATUS_INTERNAL_ERROR);
312     }
313 
314     @Test
315     public void testRunBackgroundFetchNothingToUpdate()
316             throws ExecutionException, InterruptedException {
317         assertTrue(
318                 mCustomAudienceDaoSpy
319                         .getActiveEligibleCustomAudienceBackgroundFetchData(
320                                 CommonFixture.FIXED_NOW, 1)
321                         .isEmpty());
322 
323         CountDownLatch latchForExecutionLoggerClose = new CountDownLatch(1);
324         setLatchToCountdownOnLogClose(latchForExecutionLoggerClose);
325         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
326 
327         mBackgroundFetchWorker.runBackgroundFetch().get();
328 
329         // ensures that we verify only after BackgroundExecuterLoggerClose#close is executed.
330         latchForExecutionLoggerClose.await();
331         verify(mBackgroundFetchRunnerSpy).deleteExpiredCustomAudiences(any());
332         verify(mCustomAudienceDaoSpy).deleteAllExpiredCustomAudienceData(any());
333         verify(mBackgroundFetchRunnerSpy).deleteDisallowedOwnerCustomAudiences();
334         verify(mCustomAudienceDaoSpy).deleteAllDisallowedOwnerCustomAudienceData(any(), any());
335         verify(mBackgroundFetchRunnerSpy).deleteDisallowedBuyerCustomAudiences();
336         verify(mBackgroundFetchRunnerSpy).deleteDisallowedPackageAppInstallEntries();
337         verify(mCustomAudienceDaoSpy).deleteAllDisallowedBuyerCustomAudienceData(any(), any());
338         verify(mBackgroundFetchRunnerSpy, never()).updateCustomAudience(any(), any());
339         verify(mBackgroundFetchExecutionLoggerSpy).close(0, STATUS_SUCCESS);
340     }
341 
342     @Test
343     public void testRunBackgroundFetchNothingToUpdateNoFilters()
344             throws ExecutionException, InterruptedException {
345         Flags flagsFilteringDisabled = new BackgroundFetchWorkerTestFlags(false);
346         mBackgroundFetchWorker =
347                 new BackgroundFetchWorker(
348                         mCustomAudienceDaoSpy,
349                         flagsFilteringDisabled,
350                         mBackgroundFetchRunnerSpy,
351                         mClockMock,
352                         mCustomAudienceLoggerFactoryMock);
353         assertTrue(
354                 mCustomAudienceDaoSpy
355                         .getActiveEligibleCustomAudienceBackgroundFetchData(
356                                 CommonFixture.FIXED_NOW, 1)
357                         .isEmpty());
358 
359         CountDownLatch latchForExecutionLoggerClose = new CountDownLatch(1);
360         setLatchToCountdownOnLogClose(latchForExecutionLoggerClose);
361         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
362 
363         mBackgroundFetchWorker.runBackgroundFetch().get();
364 
365         // ensures that we verify only after BackgroundExecuterLoggerClose#close is executed.
366         latchForExecutionLoggerClose.await();
367         verify(mBackgroundFetchRunnerSpy).deleteExpiredCustomAudiences(any());
368         verify(mCustomAudienceDaoSpy).deleteAllExpiredCustomAudienceData(any());
369         verify(mBackgroundFetchRunnerSpy).deleteDisallowedOwnerCustomAudiences();
370         verify(mCustomAudienceDaoSpy).deleteAllDisallowedOwnerCustomAudienceData(any(), any());
371         verify(mBackgroundFetchRunnerSpy).deleteDisallowedBuyerCustomAudiences();
372         verify(mBackgroundFetchRunnerSpy, times(0)).deleteDisallowedPackageAppInstallEntries();
373         verify(mCustomAudienceDaoSpy).deleteAllDisallowedBuyerCustomAudienceData(any(), any());
374         verify(mBackgroundFetchRunnerSpy, never()).updateCustomAudience(any(), any());
375         verify(mBackgroundFetchExecutionLoggerSpy).close(0, STATUS_SUCCESS);
376     }
377 
378     @Test
379     @FlakyTest(bugId = 298714561)
380     public void testRunBackgroundFetchUpdateOneCustomAudience()
381             throws ExecutionException, InterruptedException {
382         // Mock a single custom audience eligible for update
383         DBCustomAudienceBackgroundFetchData fetchData =
384                 DBCustomAudienceBackgroundFetchDataFixture.getValidBuilderByBuyer(
385                                 CommonFixture.VALID_BUYER_1)
386                         .setEligibleUpdateTime(CommonFixture.FIXED_NOW)
387                         .build();
388         List<DBCustomAudienceBackgroundFetchData> fetchDataList = Arrays.asList(fetchData);
389         doReturn(fetchDataList)
390                 .when(mCustomAudienceDaoSpy)
391                 .getActiveEligibleCustomAudienceBackgroundFetchData(any(), anyLong());
392         doReturn(FluentFuture.from(immediateFuture(null)))
393                 .when(mBackgroundFetchRunnerSpy)
394                 .updateCustomAudience(any(), any());
395 
396         CountDownLatch latchForExecutionLoggerClose = new CountDownLatch(1);
397         setLatchToCountdownOnLogClose(latchForExecutionLoggerClose);
398         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
399         mBackgroundFetchWorker.runBackgroundFetch().get();
400 
401         // ensures that we verify only after BackgroundExecuterLoggerClose#close is executed.
402         latchForExecutionLoggerClose.await();
403         verify(mBackgroundFetchRunnerSpy).deleteExpiredCustomAudiences(any());
404         verify(mCustomAudienceDaoSpy).deleteAllExpiredCustomAudienceData(any());
405         verify(mBackgroundFetchRunnerSpy).deleteDisallowedOwnerCustomAudiences();
406         verify(mCustomAudienceDaoSpy).deleteAllDisallowedOwnerCustomAudienceData(any(), any());
407         verify(mBackgroundFetchRunnerSpy).deleteDisallowedBuyerCustomAudiences();
408         verify(mBackgroundFetchRunnerSpy).deleteDisallowedPackageAppInstallEntries();
409         verify(mCustomAudienceDaoSpy).deleteAllDisallowedBuyerCustomAudienceData(any(), any());
410         verify(mBackgroundFetchRunnerSpy).updateCustomAudience(any(), any());
411         verify(mBackgroundFetchExecutionLoggerSpy).close(fetchDataList.size(), STATUS_SUCCESS);
412     }
413 
414     @Test
415     public void testRunBackgroundFetchUpdateCustomAudiences()
416             throws ExecutionException, InterruptedException {
417         int numEligibleCustomAudiences = 12;
418 
419         // Mock a list of custom audiences eligible for update
420         DBCustomAudienceBackgroundFetchData.Builder fetchDataBuilder =
421                 DBCustomAudienceBackgroundFetchDataFixture.getValidBuilderByBuyer(
422                                 CommonFixture.VALID_BUYER_1)
423                         .setEligibleUpdateTime(CommonFixture.FIXED_NOW);
424         List<DBCustomAudienceBackgroundFetchData> fetchDataList = new ArrayList<>();
425         for (int i = 0; i < numEligibleCustomAudiences; i++) {
426             fetchDataList.add(fetchDataBuilder.setName("ca" + i).build());
427         }
428 
429         doReturn(fetchDataList)
430                 .when(mCustomAudienceDaoSpy)
431                 .getActiveEligibleCustomAudienceBackgroundFetchData(any(), anyLong());
432         doReturn(FluentFuture.from(immediateFuture(null)))
433                 .when(mBackgroundFetchRunnerSpy)
434                 .updateCustomAudience(any(), any());
435 
436         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
437         CountDownLatch latchForExecutionLoggerClose = new CountDownLatch(1);
438         setLatchToCountdownOnLogClose(latchForExecutionLoggerClose);
439         mBackgroundFetchWorker.runBackgroundFetch().get();
440 
441         // ensures that we verify only after BackgroundExecuterLoggerClose#close is executed.
442         latchForExecutionLoggerClose.await();
443         verify(mBackgroundFetchRunnerSpy).deleteExpiredCustomAudiences(any());
444         verify(mCustomAudienceDaoSpy).deleteAllExpiredCustomAudienceData(any());
445         verify(mBackgroundFetchRunnerSpy).deleteDisallowedOwnerCustomAudiences();
446         verify(mCustomAudienceDaoSpy).deleteAllDisallowedOwnerCustomAudienceData(any(), any());
447         verify(mBackgroundFetchRunnerSpy).deleteDisallowedBuyerCustomAudiences();
448         verify(mBackgroundFetchRunnerSpy).deleteDisallowedPackageAppInstallEntries();
449         verify(mCustomAudienceDaoSpy).deleteAllDisallowedBuyerCustomAudienceData(any(), any());
450         verify(mBackgroundFetchRunnerSpy, times(numEligibleCustomAudiences))
451                 .updateCustomAudience(any(), any());
452         verify(mBackgroundFetchExecutionLoggerSpy).close(fetchDataList.size(), STATUS_SUCCESS);
453     }
454 
455     @Test
456     public void testRunBackgroundFetchChecksWorkInProgress()
457             throws InterruptedException, ExecutionException {
458         int numEligibleCustomAudiences = 16;
459         CountDownLatch partialCompletionLatch = new CountDownLatch(numEligibleCustomAudiences / 4);
460 
461         // Mock a list of custom audiences eligible for update
462         DBCustomAudienceBackgroundFetchData.Builder fetchDataBuilder =
463                 DBCustomAudienceBackgroundFetchDataFixture.getValidBuilderByBuyer(
464                                 CommonFixture.VALID_BUYER_1)
465                         .setEligibleUpdateTime(CommonFixture.FIXED_NOW);
466         List<DBCustomAudienceBackgroundFetchData> fetchDataList = new ArrayList<>();
467         for (int i = 0; i < numEligibleCustomAudiences; i++) {
468             fetchDataList.add(fetchDataBuilder.setName("ca" + i).build());
469         }
470 
471         doReturn(fetchDataList)
472                 .when(mCustomAudienceDaoSpy)
473                 .getActiveEligibleCustomAudienceBackgroundFetchData(any(), anyLong());
474         doAnswer(
475                         unusedInvocation -> {
476                             Thread.sleep(100);
477                             partialCompletionLatch.countDown();
478                             return FluentFuture.from(immediateFuture(null));
479                         })
480                 .when(mBackgroundFetchRunnerSpy)
481                 .updateCustomAudience(any(), any());
482 
483         // ensures that we verify only after BackgroundExecuterLoggerClose#close is executed.
484         CountDownLatch latchForExecutionLoggerClose = new CountDownLatch(1);
485         setLatchToCountdownOnLogClose(latchForExecutionLoggerClose);
486         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
487 
488         CountDownLatch bgfWorkStoppedLatch = new CountDownLatch(1);
489         mExecutorService.execute(
490                 () -> {
491                     try {
492                         mBackgroundFetchWorker.runBackgroundFetch().get();
493                     } catch (Exception exception) {
494                         sLogger.e(
495                                 exception, "Exception encountered while running background fetch");
496                     } finally {
497                         bgfWorkStoppedLatch.countDown();
498                     }
499                 });
500 
501         // Wait til updates are partially complete, then try running background fetch again and
502         // verify nothing is done
503         partialCompletionLatch.await();
504         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW.plusSeconds(1));
505         mBackgroundFetchWorker.runBackgroundFetch().get();
506 
507         bgfWorkStoppedLatch.await();
508         latchForExecutionLoggerClose.await();
509         verify(mBackgroundFetchRunnerSpy).deleteExpiredCustomAudiences(any());
510         verify(mCustomAudienceDaoSpy).deleteAllExpiredCustomAudienceData(any());
511         verify(mBackgroundFetchRunnerSpy).deleteDisallowedOwnerCustomAudiences();
512         verify(mCustomAudienceDaoSpy).deleteAllDisallowedOwnerCustomAudienceData(any(), any());
513         verify(mBackgroundFetchRunnerSpy).deleteDisallowedBuyerCustomAudiences();
514         verify(mBackgroundFetchRunnerSpy).deleteDisallowedPackageAppInstallEntries();
515         verify(mCustomAudienceDaoSpy).deleteAllDisallowedBuyerCustomAudienceData(any(), any());
516         verify(mBackgroundFetchRunnerSpy, times(numEligibleCustomAudiences))
517                 .updateCustomAudience(any(), any());
518         verify(mBackgroundFetchExecutionLoggerSpy).close(fetchDataList.size(), STATUS_SUCCESS);
519     }
520 
521     @Test
522     public void testStopWorkWithoutRunningFetchDoesNothing() {
523         // Verify no errors/exceptions thrown when no work in progress
524         mBackgroundFetchWorker.stopWork();
525     }
526 
527     @Test
528     public void testStopWorkGracefullyStopsBackgroundFetch() throws Exception {
529         int numEligibleCustomAudiences = 16;
530         CountDownLatch partialCompletionLatch = new CountDownLatch(numEligibleCustomAudiences / 4);
531 
532         // Mock a list of custom audiences eligible for update
533         DBCustomAudienceBackgroundFetchData.Builder fetchDataBuilder =
534                 DBCustomAudienceBackgroundFetchDataFixture.getValidBuilderByBuyer(
535                                 CommonFixture.VALID_BUYER_1)
536                         .setEligibleUpdateTime(CommonFixture.FIXED_NOW);
537         List<DBCustomAudienceBackgroundFetchData> fetchDataList = new ArrayList<>();
538         for (int i = 0; i < numEligibleCustomAudiences; i++) {
539             fetchDataList.add(fetchDataBuilder.setName("ca" + i).build());
540         }
541 
542         doReturn(fetchDataList)
543                 .when(mCustomAudienceDaoSpy)
544                 .getActiveEligibleCustomAudienceBackgroundFetchData(any(), anyLong());
545         doAnswer(
546                         unusedInvocation -> {
547                             Thread.sleep(100);
548                             partialCompletionLatch.countDown();
549                             return FluentFuture.from(immediateVoidFuture());
550                         })
551                 .when(mBackgroundFetchRunnerSpy)
552                 .updateCustomAudience(any(), any());
553 
554         // Ensure that logger is closed after stopping work.
555         CountDownLatch closeLoggerLatch = new CountDownLatch(1);
556         doAnswer(
557                         unusedInovcation -> {
558                             closeLoggerLatch.countDown();
559                             return null;
560                         })
561                 .when(mBackgroundFetchExecutionLoggerSpy)
562                 .close(anyInt(), anyInt());
563 
564         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
565 
566         ListenableFuture<Void> backgroundFetchResult = mBackgroundFetchWorker.runBackgroundFetch();
567 
568         // Wait til updates are partially complete, then try stopping background fetch
569         partialCompletionLatch.await();
570         mBackgroundFetchWorker.stopWork();
571         // stopWork() should notify to the worker that the work should end so the future
572         // should complete within the time required to update the custom audiences
573         backgroundFetchResult.get(
574                 100 * (numEligibleCustomAudiences * 3 / 4) + 100, TimeUnit.SECONDS);
575         // Wait for logger to close.
576         assertThat(closeLoggerLatch.await(200, TimeUnit.MILLISECONDS)).isTrue();
577         verify(mBackgroundFetchExecutionLoggerSpy)
578                 .close(numEligibleCustomAudiences, STATUS_SUCCESS);
579     }
580 
581     @Test
582     @FlakyTest(bugId = 316155251)
583     public void testStopWorkPreemptsDataUpdates() throws Exception {
584         int numEligibleCustomAudiences = 16;
585         CountDownLatch beforeUpdatingCasLatch = new CountDownLatch(numEligibleCustomAudiences / 4);
586 
587         // Mock a list of custom audiences eligible for update
588         DBCustomAudienceBackgroundFetchData.Builder fetchDataBuilder =
589                 DBCustomAudienceBackgroundFetchDataFixture.getValidBuilderByBuyer(
590                                 CommonFixture.VALID_BUYER_1)
591                         .setEligibleUpdateTime(CommonFixture.FIXED_NOW);
592         List<DBCustomAudienceBackgroundFetchData> fetchDataList = new ArrayList<>();
593         for (int i = 0; i < numEligibleCustomAudiences; i++) {
594             fetchDataList.add(fetchDataBuilder.setName("ca" + i).build());
595         }
596 
597         // Ensuring that stopWork is called before the data update process
598         doAnswer(
599                         unusedInvocation -> {
600                             beforeUpdatingCasLatch.await();
601                             return fetchDataList;
602                         })
603                 .when(mCustomAudienceDaoSpy)
604                 .getActiveEligibleCustomAudienceBackgroundFetchData(any(), anyLong());
605         doAnswer(
606                         unusedInvocation -> {
607                             Thread.sleep(100);
608                             return null;
609                         })
610                 .when(mBackgroundFetchRunnerSpy)
611                 .updateCustomAudience(any(), any());
612 
613         // Ensure that logger is closed after stopping work.
614         CountDownLatch closeLoggerLatch = new CountDownLatch(1);
615         doAnswer(
616                         unusedInovcation -> {
617                             closeLoggerLatch.countDown();
618                             return null;
619                         })
620                 .when(mBackgroundFetchExecutionLoggerSpy)
621                 .close(anyInt(), anyInt());
622 
623         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
624 
625         ListenableFuture<Void> backgroundFetchResult = mBackgroundFetchWorker.runBackgroundFetch();
626 
627         // Wait til updates are partially complete, then try stopping background fetch
628         mBackgroundFetchWorker.stopWork();
629         beforeUpdatingCasLatch.countDown();
630         // stopWork() called before updating the data should cause immediate termination
631         // waiting for 200ms to handle thread scheduling delays.
632         // The important check is that the time is less than the time of updating all CAs
633         backgroundFetchResult.get(200, TimeUnit.MILLISECONDS);
634         // Wait for logger to close.
635         assertThat(closeLoggerLatch.await(200, TimeUnit.MILLISECONDS)).isTrue();
636         verify(mBackgroundFetchExecutionLoggerSpy).close(0, STATUS_SUCCESS);
637     }
638 
639     @Test
640     public void testRunBackgroundFetchInSequence() throws InterruptedException, ExecutionException {
641         int numEligibleCustomAudiences = 16;
642         int expectedUpdateCustomAudienceCalls = numEligibleCustomAudiences / 2;
643 
644         // Mock two lists of custom audiences eligible for update
645         DBCustomAudienceBackgroundFetchData.Builder fetchDataBuilder =
646                 DBCustomAudienceBackgroundFetchDataFixture.getValidBuilderByBuyer(
647                                 CommonFixture.VALID_BUYER_1)
648                         .setEligibleUpdateTime(CommonFixture.FIXED_NOW);
649         List<DBCustomAudienceBackgroundFetchData> fetchDataList1 = new ArrayList<>();
650         List<DBCustomAudienceBackgroundFetchData> fetchDataList2 = new ArrayList<>();
651         for (int i = 0; i < numEligibleCustomAudiences; i++) {
652             DBCustomAudienceBackgroundFetchData fetchData =
653                     fetchDataBuilder.setName("ca" + i).build();
654             if (i < expectedUpdateCustomAudienceCalls) {
655                 fetchDataList1.add(fetchData);
656             } else {
657                 fetchDataList2.add(fetchData);
658             }
659         }
660 
661         // Count the number of times updateCustomAudience is run
662         AtomicInteger completionCount = new AtomicInteger(0);
663 
664         // Return the first list the first time, and the second list in the second call
665         AnswerSyncCallback<FluentFuture<UpdateResultType>> updateCustomAudienceCallback =
666                 AnswerSyncCallback.forMultipleAnswers(
667                         FluentFuture.from(immediateFuture(null)),
668                         expectedUpdateCustomAudienceCalls);
669 
670         doReturn(fetchDataList1)
671                 .doReturn(fetchDataList2)
672                 .when(mCustomAudienceDaoSpy)
673                 .getActiveEligibleCustomAudienceBackgroundFetchData(any(), anyLong());
674         doAnswer(updateCustomAudienceCallback)
675                 .when(mBackgroundFetchRunnerSpy)
676                 .updateCustomAudience(any(), any());
677 
678         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
679 
680         // TODO(b/321743128): Clarify the production behavior for the logger and use existed mocking
681         // method setLatchToCountdownOnLogClose().
682         //
683         // Verify the invocations of the logging events.
684         AnswerSyncCallback<Void> loggerCloseCallback =
685                 AnswerSyncCallback.forMultipleVoidAnswers(/* numberOfExpectedCalls= */ 2);
686         doAnswer(loggerCloseCallback)
687                 .when(mBackgroundFetchExecutionLoggerSpy)
688                 .close(anyInt(), anyInt());
689 
690         SimpleSyncCallback bgfWorkStoppedCallback = new SimpleSyncCallback();
691         mExecutorService.execute(
692                 () -> {
693                     try {
694                         mBackgroundFetchWorker.runBackgroundFetch().get();
695                     } catch (Exception exception) {
696                         sLogger.e(
697                                 exception, "Exception encountered while running background fetch");
698                     } finally {
699                         bgfWorkStoppedCallback.setCalled();
700                     }
701                 });
702 
703         // Wait til updates are complete, then try running background fetch again and
704         // verify the second run updates more custom audiences successfully
705         updateCustomAudienceCallback.assertCalled();
706         bgfWorkStoppedCallback.assertCalled();
707         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW.plusSeconds(1));
708         mBackgroundFetchWorker.runBackgroundFetch().get();
709 
710         verify(mBackgroundFetchRunnerSpy, times(2)).deleteExpiredCustomAudiences(any());
711         verify(mCustomAudienceDaoSpy, times(2)).deleteAllExpiredCustomAudienceData(any());
712         verify(mBackgroundFetchRunnerSpy, times(2)).deleteDisallowedOwnerCustomAudiences();
713         verify(mCustomAudienceDaoSpy, times(2))
714                 .deleteAllDisallowedOwnerCustomAudienceData(any(), any());
715         verify(mBackgroundFetchRunnerSpy, times(2)).deleteDisallowedBuyerCustomAudiences();
716         verify(mBackgroundFetchRunnerSpy, times(2)).deleteDisallowedPackageAppInstallEntries();
717         verify(mCustomAudienceDaoSpy, times(2))
718                 .deleteAllDisallowedBuyerCustomAudienceData(any(), any());
719         verify(mBackgroundFetchRunnerSpy, times(numEligibleCustomAudiences))
720                 .updateCustomAudience(any(), any());
721         assertThat(updateCustomAudienceCallback.getNumberActualCalls())
722                 .isEqualTo(numEligibleCustomAudiences);
723         loggerCloseCallback.assertCalled();
724     }
725 
726     @Test
727     public void testRunBackgroundFetchHandlesLoggerCloseError()
728             throws ExecutionException, InterruptedException {
729         // Mock a single custom audience eligible for update
730         DBCustomAudienceBackgroundFetchData fetchData =
731                 DBCustomAudienceBackgroundFetchDataFixture.getValidBuilderByBuyer(
732                                 CommonFixture.VALID_BUYER_1)
733                         .setEligibleUpdateTime(CommonFixture.FIXED_NOW)
734                         .build();
735         List<DBCustomAudienceBackgroundFetchData> fetchDataList = Arrays.asList(fetchData);
736         doReturn(fetchDataList)
737                 .when(mCustomAudienceDaoSpy)
738                 .getActiveEligibleCustomAudienceBackgroundFetchData(any(), anyLong());
739         doReturn(FluentFuture.from(immediateFuture(null)))
740                 .when(mBackgroundFetchRunnerSpy)
741                 .updateCustomAudience(any(), any());
742 
743         when(mClockMock.instant()).thenReturn(CommonFixture.FIXED_NOW);
744 
745         doThrow(new IllegalStateException())
746                 .when(mBackgroundFetchExecutionLoggerSpy)
747                 .close(anyInt(), anyInt());
748 
749         // Background fetch should complete without error
750         mBackgroundFetchWorker.runBackgroundFetch().get();
751     }
752 
753     private void setLatchToCountdownOnLogClose(CountDownLatch latch) {
754         doAnswer(
755                         unusedInvocation -> {
756                             latch.countDown();
757                             return null;
758                         })
759                 .when(mAdServicesLoggerImplMock)
760                 .logBackgroundFetchProcessReportedStats(any());
761     }
762 
763     private static class BackgroundFetchWorkerTestFlags implements Flags {
764         private final boolean mFledgeAppInstallFilteringEnabled;
765 
766         BackgroundFetchWorkerTestFlags(boolean fledgeAppInstallFilteringEnabled) {
767             mFledgeAppInstallFilteringEnabled = fledgeAppInstallFilteringEnabled;
768         }
769 
770         @Override
771         public int getFledgeBackgroundFetchThreadPoolSize() {
772             return 4;
773         }
774 
775         @Override
776         public boolean getFledgeAppInstallFilteringEnabled() {
777             return mFledgeAppInstallFilteringEnabled;
778         }
779 
780         @Override
781         public int getFledgeBackgroundFetchNetworkConnectTimeoutMs() {
782             return EXTENDED_FLEDGE_BACKGROUND_FETCH_NETWORK_CONNECT_TIMEOUT_MS;
783         }
784 
785         @Override
786         public int getFledgeBackgroundFetchNetworkReadTimeoutMs() {
787             return EXTENDED_FLEDGE_BACKGROUND_FETCH_NETWORK_READ_TIMEOUT_MS;
788         }
789     }
790 }
791