/* * Copyright (C) 2018 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.textclassifier.downloader; import static com.google.common.truth.Truth.assertThat; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import android.content.Context; import android.os.LocaleList; import androidx.room.Room; import androidx.test.core.app.ApplicationProvider; import androidx.work.ListenableWorker; import androidx.work.WorkerFactory; import androidx.work.WorkerParameters; import androidx.work.testing.TestListenableWorkerBuilder; import com.android.os.AtomsProto.TextClassifierDownloadReported; import com.android.os.AtomsProto.TextClassifierDownloadReported.DownloadStatus; import com.android.os.AtomsProto.TextClassifierDownloadReported.FailureReason; import com.android.os.AtomsProto.TextClassifierDownloadWorkCompleted; import com.android.os.AtomsProto.TextClassifierDownloadWorkCompleted.WorkResult; import com.android.textclassifier.common.ModelType; import com.android.textclassifier.common.TextClassifierSettings; import com.android.textclassifier.common.statsd.TextClassifierDownloadLoggerTestRule; import com.android.textclassifier.testing.SetDefaultLocalesRule; import com.android.textclassifier.testing.TestingDeviceConfig; import com.google.common.collect.Iterables; import com.google.common.util.concurrent.Futures; import com.google.common.util.concurrent.MoreExecutors; import java.io.File; import java.time.Clock; import java.time.Instant; import java.util.List; import java.util.Locale; import java.util.stream.Collectors; import org.junit.After; import org.junit.Before; import org.junit.Rule; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; import org.mockito.Mock; import org.mockito.MockitoAnnotations; @RunWith(JUnit4.class) public final class ModelDownloadWorkerTest { private static final long WORK_ID = 123456789L; private static final String MODEL_TYPE = ModelType.ANNOTATOR; private static final String MODEL_TYPE_2 = ModelType.ACTIONS_SUGGESTIONS; private static final TextClassifierDownloadReported.ModelType MODEL_TYPE_ATOM = TextClassifierDownloadReported.ModelType.ANNOTATOR; private static final String LOCALE_TAG = "en"; private static final String LOCALE_TAG_2 = "zh"; private static final String LOCALE_TAG_3 = "es"; private static final String MANIFEST_URL = "https://www.gstatic.com/android/text_classifier/q/v711/en.fb.manifest"; private static final String MANIFEST_URL_2 = "https://www.gstatic.com/android/text_classifier/q/v711/zh.fb.manifest"; private static final String MANIFEST_URL_3 = "https://www.gstatic.com/android/text_classifier/q/v711/es.fb.manifest"; private static final String MODEL_URL = "https://www.gstatic.com/android/text_classifier/q/v711/en.fb"; private static final String MODEL_URL_2 = "https://www.gstatic.com/android/text_classifier/q/v711/zh.fb"; private static final String MODEL_URL_3 = "https://www.gstatic.com/android/text_classifier/q/v711/es.fb"; private static final int RUN_ATTEMPT_COUNT = 1; private static final int WORKER_MAX_RUN_ATTEMPT_COUNT = 5; private static final int MANIFEST_MAX_ATTEMPT_COUNT = 2; private static final ModelManifest.Model MODEL_PROTO = ModelManifest.Model.newBuilder() .setUrl(MODEL_URL) .setSizeInBytes(1) .setFingerprint("fingerprint") .build(); private static final ModelManifest.Model MODEL_PROTO_2 = ModelManifest.Model.newBuilder() .setUrl(MODEL_URL_2) .setSizeInBytes(1) .setFingerprint("fingerprint") .build(); private static final ModelManifest.Model MODEL_PROTO_3 = ModelManifest.Model.newBuilder() .setUrl(MODEL_URL_3) .setSizeInBytes(1) .setFingerprint("fingerprint") .build(); private static final ModelManifest MODEL_MANIFEST_PROTO = ModelManifest.newBuilder().addModels(MODEL_PROTO).build(); private static final ModelManifest MODEL_MANIFEST_PROTO_2 = ModelManifest.newBuilder().addModels(MODEL_PROTO_2).build(); private static final ModelManifest MODEL_MANIFEST_PROTO_3 = ModelManifest.newBuilder().addModels(MODEL_PROTO_3).build(); private static final ModelDownloadException FAILED_TO_DOWNLOAD_EXCEPTION = new ModelDownloadException( ModelDownloadException.FAILED_TO_DOWNLOAD_OTHER, "failed to download"); private static final FailureReason FAILED_TO_DOWNLOAD_FAILURE_REASON = TextClassifierDownloadReported.FailureReason.FAILED_TO_DOWNLOAD_OTHER; private static final LocaleList DEFAULT_LOCALE_LIST = new LocaleList(new Locale(LOCALE_TAG)); private static final LocaleList LOCALE_LIST_2 = new LocaleList(new Locale(LOCALE_TAG), new Locale(LOCALE_TAG_2)); private static final LocaleList LOCALE_LIST_3 = new LocaleList(new Locale(LOCALE_TAG), new Locale(LOCALE_TAG_2), new Locale(LOCALE_TAG_3)); private static final Instant WORK_SCHEDULED_TIME = Instant.now(); private static final Instant WORK_STARTED_TIME = WORK_SCHEDULED_TIME.plusSeconds(100); // Make sure any combination has a different diff private static final Instant DOWNLOAD_STARTED_TIME = WORK_STARTED_TIME.plusSeconds(1); private static final Instant DOWNLOAD_ENDED_TIME = WORK_STARTED_TIME.plusSeconds(1 + 2); private static final Instant DOWNLOAD_STARTED_TIME_2 = WORK_STARTED_TIME.plusSeconds(1 + 2 + 3); private static final Instant DOWNLOAD_ENDED_TIME_2 = WORK_STARTED_TIME.plusSeconds(1 + 2 + 3 + 4); private static final Instant WORK_ENDED_TIME = WORK_STARTED_TIME.plusSeconds(1 + 2 + 3 + 4 + 5); private static final long DOWNLOAD_STARTED_TO_ENDED_MILLIS = DOWNLOAD_ENDED_TIME.toEpochMilli() - DOWNLOAD_STARTED_TIME.toEpochMilli(); private static final long DOWNLOAD_STARTED_TO_ENDED_2_MILLIS = DOWNLOAD_ENDED_TIME_2.toEpochMilli() - DOWNLOAD_STARTED_TIME_2.toEpochMilli(); private static final long WORK_SCHEDULED_TO_STARTED_MILLIS = WORK_STARTED_TIME.toEpochMilli() - WORK_SCHEDULED_TIME.toEpochMilli(); private static final long WORK_STARTED_TO_ENDED_MILLIS = WORK_ENDED_TIME.toEpochMilli() - WORK_STARTED_TIME.toEpochMilli(); @Mock private Clock clock; @Mock private ModelDownloader modelDownloader; private File modelDownloaderDir; private File modelFile; private File modelFile2; private File modelFile3; private DownloadedModelDatabase db; private DownloadedModelManager downloadedModelManager; private TestingDeviceConfig deviceConfig; private TextClassifierSettings settings; @Rule public final SetDefaultLocalesRule setDefaultLocalesRule = new SetDefaultLocalesRule(); @Rule public final TextClassifierDownloadLoggerTestRule loggerTestRule = new TextClassifierDownloadLoggerTestRule(); @Before public void setUp() { MockitoAnnotations.initMocks(this); Context context = ApplicationProvider.getApplicationContext(); this.deviceConfig = new TestingDeviceConfig(); this.settings = new TextClassifierSettings(deviceConfig, /* isWear= */ false); this.modelDownloaderDir = new File(context.getCacheDir(), "downloaded"); this.modelDownloaderDir.mkdirs(); this.modelFile = new File(modelDownloaderDir, "test.model"); this.modelFile2 = new File(modelDownloaderDir, "test2.model"); this.modelFile3 = new File(modelDownloaderDir, "test3.model"); this.db = Room.inMemoryDatabaseBuilder(context, DownloadedModelDatabase.class).build(); this.downloadedModelManager = DownloadedModelManagerImpl.getInstanceForTesting(db, modelDownloaderDir, settings); setDefaultLocalesRule.set(DEFAULT_LOCALE_LIST); deviceConfig.setConfig(TextClassifierSettings.MODEL_DOWNLOAD_MANAGER_ENABLED, true); } @After public void cleanUp() { db.close(); DownloaderTestUtils.deleteRecursively(modelDownloaderDir); } @Test public void downloadSucceed() throws Exception { setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); modelFile.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); verifySucceededDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceed_modelAlreadyExists() throws Exception { setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); modelFile.createNewFile(); downloadedModelManager.registerModel(MODEL_URL, modelFile.getAbsolutePath()); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); verifySucceededDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceed_manifestAlreadyExists() throws Exception { setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); modelFile.createNewFile(); downloadedModelManager.registerModel(MODEL_URL, modelFile.getAbsolutePath()); downloadedModelManager.registerManifest(MANIFEST_URL, MODEL_URL); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); verifySucceededDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceed_downloadMultipleModels() throws Exception { setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE_2, LOCALE_TAG, MANIFEST_URL_2); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadManifest(MANIFEST_URL_2)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO_2)); modelFile.createNewFile(); modelFile2.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO_2)) .thenReturn(Futures.immediateFuture(modelFile2)); // We assume we always download MODEL_TYPE first and then MODEL_TYPE_2, o/w this will be flaky when(clock.instant()) .thenReturn( WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, DOWNLOAD_STARTED_TIME_2, DOWNLOAD_ENDED_TIME_2, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(modelFile2.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); assertThat(downloadedModelManager.listModels(MODEL_TYPE_2)).containsExactly(modelFile2); List atoms = loggerTestRule.getLoggedDownloadReportedAtoms(); assertThat(atoms).hasSize(2); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getUrlSuffix) .collect(Collectors.toList())) .containsExactly(MANIFEST_URL, MANIFEST_URL_2); assertThat(atoms.get(0).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat(atoms.get(1).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getDownloadDurationMillis) .collect(Collectors.toList())) .containsExactly(DOWNLOAD_STARTED_TO_ENDED_MILLIS, DOWNLOAD_STARTED_TO_ENDED_2_MILLIS); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceed_shareSingleModelDownloadForMultipleManifest() throws Exception { setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE_2, LOCALE_TAG, MANIFEST_URL_2); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadManifest(MANIFEST_URL_2)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); modelFile.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); // We assume we always download MODEL_TYPE first and then MODEL_TYPE_2, o/w this will be flaky when(clock.instant()) .thenReturn( WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, DOWNLOAD_STARTED_TIME_2, DOWNLOAD_ENDED_TIME_2, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); assertThat(downloadedModelManager.listModels(MODEL_TYPE_2)).containsExactly(modelFile); verify(modelDownloader, times(1)).downloadModel(modelDownloaderDir, MODEL_PROTO); List atoms = loggerTestRule.getLoggedDownloadReportedAtoms(); assertThat(atoms).hasSize(2); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getUrlSuffix) .collect(Collectors.toList())) .containsExactly(MANIFEST_URL, MANIFEST_URL_2); assertThat(atoms.get(0).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat(atoms.get(1).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getDownloadDurationMillis) .collect(Collectors.toList())) .containsExactly(DOWNLOAD_STARTED_TO_ENDED_MILLIS, DOWNLOAD_STARTED_TO_ENDED_2_MILLIS); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceed_shareManifest() throws Exception { setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE_2, LOCALE_TAG, MANIFEST_URL); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); modelFile.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); // We assume we always download MODEL_TYPE first and then MODEL_TYPE_2, o/w this will be flaky when(clock.instant()) .thenReturn( WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, DOWNLOAD_STARTED_TIME_2, DOWNLOAD_ENDED_TIME_2, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); assertThat(downloadedModelManager.listModels(MODEL_TYPE_2)).containsExactly(modelFile); verify(modelDownloader, times(1)).downloadManifest(MANIFEST_URL); List atoms = loggerTestRule.getLoggedDownloadReportedAtoms(); assertThat(atoms).hasSize(2); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getUrlSuffix) .collect(Collectors.toList())) .containsExactly(MANIFEST_URL, MANIFEST_URL); assertThat(atoms.get(0).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat(atoms.get(1).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getDownloadDurationMillis) .collect(Collectors.toList())) .containsExactly(DOWNLOAD_STARTED_TO_ENDED_MILLIS, DOWNLOAD_STARTED_TO_ENDED_2_MILLIS); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadFailed_failedToDownloadManifest() throws Exception { setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFailedFuture(FAILED_TO_DOWNLOAD_EXCEPTION)); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.retry()); verifyFailedDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.RETRY_MODEL_DOWNLOAD_FAILED); } @Test public void downloadFailed_failedToDownloadModel() throws Exception { setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFailedFuture(FAILED_TO_DOWNLOAD_EXCEPTION)); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.retry()); verifyFailedDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.RETRY_MODEL_DOWNLOAD_FAILED); } @Test public void downloadFailed_modelDownloadManagerDisabled() throws Exception { deviceConfig.setConfig(TextClassifierSettings.MODEL_DOWNLOAD_MANAGER_ENABLED, false); when(clock.instant()).thenReturn(WORK_STARTED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.failure()); assertThat(loggerTestRule.getLoggedDownloadReportedAtoms()).isEmpty(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.FAILURE_MODEL_DOWNLOADER_DISABLED); } @Test public void downloadFailed_reachWorkerMaxRunAttempts() throws Exception { when(clock.instant()).thenReturn(WORK_STARTED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(WORKER_MAX_RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.failure()); assertThat(loggerTestRule.getLoggedDownloadReportedAtoms()).isEmpty(); verifyWorkLogging(WORKER_MAX_RUN_ATTEMPT_COUNT, WorkResult.FAILURE_MAX_RUN_ATTEMPT_REACHED); } @Test public void downloadSkipped_reachManifestMaxAttempts() throws Exception { when(clock.instant()).thenReturn(WORK_STARTED_TIME, WORK_ENDED_TIME); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); deviceConfig.setConfig( TextClassifierSettings.MANIFEST_DOWNLOAD_MAX_ATTEMPTS, MANIFEST_MAX_ATTEMPT_COUNT); for (int i = 0; i < MANIFEST_MAX_ATTEMPT_COUNT; i++) { downloadedModelManager.registerManifestDownloadFailure(MANIFEST_URL); } ModelDownloadWorker worker = createWorker(MANIFEST_MAX_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(loggerTestRule.getLoggedDownloadReportedAtoms()).isEmpty(); verifyWorkLogging(MANIFEST_MAX_ATTEMPT_COUNT, WorkResult.SUCCESS_NO_UPDATE_AVAILABLE); } @Test public void downloadSkipped_manifestAlreadyProcessed() throws Exception { when(clock.instant()).thenReturn(WORK_STARTED_TIME, WORK_ENDED_TIME); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); modelFile.createNewFile(); downloadedModelManager.registerModel(MODEL_URL, modelFile.getAbsolutePath()); downloadedModelManager.registerManifest(MANIFEST_URL, MODEL_URL); downloadedModelManager.registerManifestEnrollment(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(loggerTestRule.getLoggedDownloadReportedAtoms()).isEmpty(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_NO_UPDATE_AVAILABLE); } @Test public void downloadSucceeded_multiLanguageSupportEnabled() throws Exception { setDefaultLocalesRule.set(LOCALE_LIST_2); deviceConfig.setConfig(TextClassifierSettings.MULTI_LANGUAGE_SUPPORT_ENABLED, true); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG_2, MANIFEST_URL_2); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadManifest(MANIFEST_URL_2)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO_2)); modelFile.createNewFile(); modelFile2.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO_2)) .thenReturn(Futures.immediateFuture(modelFile2)); // We assume we always download MODEL_TYPE first and then MODEL_TYPE_2, o/w this will be flaky when(clock.instant()) .thenReturn( WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, DOWNLOAD_STARTED_TIME_2, DOWNLOAD_ENDED_TIME_2, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(modelFile2.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)) .containsExactly(modelFile, modelFile2); List atoms = loggerTestRule.getLoggedDownloadReportedAtoms(); assertThat(atoms).hasSize(2); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getUrlSuffix) .collect(Collectors.toList())) .containsExactly(MANIFEST_URL, MANIFEST_URL_2); assertThat(atoms.get(0).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat(atoms.get(1).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getDownloadDurationMillis) .collect(Collectors.toList())) .containsExactly(DOWNLOAD_STARTED_TO_ENDED_MILLIS, DOWNLOAD_STARTED_TO_ENDED_2_MILLIS); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceeded_multiLanguageSupportEnabled_singleLocale() throws Exception { setDefaultLocalesRule.set(DEFAULT_LOCALE_LIST); deviceConfig.setConfig(TextClassifierSettings.MULTI_LANGUAGE_SUPPORT_ENABLED, true); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); modelFile.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); verifySucceededDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceeded_multiLanguageSupportEnabled_oneManifestAlreadyDownloaded() throws Exception { setDefaultLocalesRule.set(LOCALE_LIST_2); deviceConfig.setConfig(TextClassifierSettings.MULTI_LANGUAGE_SUPPORT_ENABLED, true); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG_2, MANIFEST_URL_2); modelFile.createNewFile(); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); modelFile2.createNewFile(); downloadedModelManager.registerModel(MODEL_URL_2, modelFile2.getAbsolutePath()); downloadedModelManager.registerManifest(MANIFEST_URL_2, MODEL_URL_2); downloadedModelManager.registerManifestEnrollment(MODEL_TYPE, LOCALE_TAG_2, MANIFEST_URL_2); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(modelFile2.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)) .containsExactly(modelFile, modelFile2); verifySucceededDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceeded_multiLanguageSupportDisabled() throws Exception { setDefaultLocalesRule.set(LOCALE_LIST_2); deviceConfig.setConfig(TextClassifierSettings.MULTI_LANGUAGE_SUPPORT_ENABLED, false); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG_2, MANIFEST_URL_2); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadManifest(MANIFEST_URL_2)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO_2)); modelFile.createNewFile(); modelFile2.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO_2)) .thenReturn(Futures.immediateFuture(modelFile2)); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); verifySucceededDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void oneDownloadFailed_multiLanguageSupportEnabled() throws Exception { setDefaultLocalesRule.set(LOCALE_LIST_2); deviceConfig.setConfig(TextClassifierSettings.MULTI_LANGUAGE_SUPPORT_ENABLED, true); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG_2, MANIFEST_URL_2); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadManifest(MANIFEST_URL_2)) .thenReturn(Futures.immediateFailedFuture(FAILED_TO_DOWNLOAD_EXCEPTION)); modelFile.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); // We assume we always download MODEL_TYPE first and then MODEL_TYPE_2, o/w this will be flaky when(clock.instant()) .thenReturn( WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, DOWNLOAD_STARTED_TIME_2, DOWNLOAD_ENDED_TIME_2, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.retry()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); List atoms = loggerTestRule.getLoggedDownloadReportedAtoms(); assertThat(atoms).hasSize(2); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getUrlSuffix) .collect(Collectors.toList())) .containsExactly(MANIFEST_URL, MANIFEST_URL_2); assertThat(atoms.get(0).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat(atoms.get(1).getDownloadStatus()).isEqualTo(DownloadStatus.FAILED_AND_RETRY); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getDownloadDurationMillis) .collect(Collectors.toList())) .containsExactly(DOWNLOAD_STARTED_TO_ENDED_MILLIS, DOWNLOAD_STARTED_TO_ENDED_2_MILLIS); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.RETRY_MODEL_DOWNLOAD_FAILED); } @Test public void downloadSucceeded_multiLanguageSupportEnabled_checkLimit() throws Exception { setDefaultLocalesRule.set(LOCALE_LIST_3); deviceConfig.setConfig(TextClassifierSettings.MULTI_LANGUAGE_SUPPORT_ENABLED, true); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG_2, MANIFEST_URL_2); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG_3, MANIFEST_URL_3); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadManifest(MANIFEST_URL_2)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO_2)); when(modelDownloader.downloadManifest(MANIFEST_URL_3)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO_3)); modelFile.createNewFile(); modelFile2.createNewFile(); modelFile3.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO_2)) .thenReturn(Futures.immediateFuture(modelFile2)); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO_3)) .thenReturn(Futures.immediateFuture(modelFile3)); // We assume we always download MODEL_TYPE first and then MODEL_TYPE_2, o/w this will be flaky when(clock.instant()) .thenReturn( WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, DOWNLOAD_STARTED_TIME_2, DOWNLOAD_ENDED_TIME_2, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(modelFile2.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)) .containsExactly(modelFile, modelFile2); List atoms = loggerTestRule.getLoggedDownloadReportedAtoms(); assertThat(atoms).hasSize(2); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getUrlSuffix) .collect(Collectors.toList())) .containsExactly(MANIFEST_URL, MANIFEST_URL_2); assertThat(atoms.get(0).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat(atoms.get(1).getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat( atoms.stream() .map(TextClassifierDownloadReported::getDownloadDurationMillis) .collect(Collectors.toList())) .containsExactly(DOWNLOAD_STARTED_TO_ENDED_MILLIS, DOWNLOAD_STARTED_TO_ENDED_2_MILLIS); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } @Test public void downloadSucceed_multiLanguageSupportEnabled_onlyDownloadMultipleForEnabledModelType() throws Exception { setDefaultLocalesRule.set(LOCALE_LIST_2); deviceConfig.setConfig(TextClassifierSettings.MULTI_LANGUAGE_SUPPORT_ENABLED, true); deviceConfig.setConfig( TextClassifierSettings.ENABLED_MODEL_TYPES_FOR_MULTI_LANGUAGE_SUPPORT, MODEL_TYPE_2); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG, MANIFEST_URL); setUpManifestUrl(MODEL_TYPE, LOCALE_TAG_2, MANIFEST_URL_2); when(modelDownloader.downloadManifest(MANIFEST_URL)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO)); when(modelDownloader.downloadManifest(MANIFEST_URL_2)) .thenReturn(Futures.immediateFuture(MODEL_MANIFEST_PROTO_2)); modelFile.createNewFile(); modelFile2.createNewFile(); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO)) .thenReturn(Futures.immediateFuture(modelFile)); when(modelDownloader.downloadModel(modelDownloaderDir, MODEL_PROTO_2)) .thenReturn(Futures.immediateFuture(modelFile2)); when(clock.instant()) .thenReturn(WORK_STARTED_TIME, DOWNLOAD_STARTED_TIME, DOWNLOAD_ENDED_TIME, WORK_ENDED_TIME); ModelDownloadWorker worker = createWorker(RUN_ATTEMPT_COUNT); assertThat(worker.startWork().get()).isEqualTo(ListenableWorker.Result.success()); assertThat(modelFile.exists()).isTrue(); assertThat(downloadedModelManager.listModels(MODEL_TYPE)).containsExactly(modelFile); verifySucceededDownloadLogging(); verifyWorkLogging(RUN_ATTEMPT_COUNT, WorkResult.SUCCESS_MODEL_DOWNLOADED); } private ModelDownloadWorker createWorker(int runAttemptCount) { return TestListenableWorkerBuilder.from( ApplicationProvider.getApplicationContext(), ModelDownloadWorker.class) .setRunAttemptCount(runAttemptCount) .setWorkerFactory( new WorkerFactory() { @Override public ListenableWorker createWorker( Context appContext, String workerClassName, WorkerParameters workerParameters) { return new ModelDownloadWorker( appContext, workerParameters, MoreExecutors.newDirectExecutorService(), modelDownloader, downloadedModelManager, settings, WORK_ID, clock, WORK_SCHEDULED_TIME.toEpochMilli()); } }) .build(); } private void verifySucceededDownloadLogging() throws Exception { TextClassifierDownloadReported atom = Iterables.getOnlyElement(loggerTestRule.getLoggedDownloadReportedAtoms()); assertThat(atom.getWorkId()).isEqualTo(WORK_ID); assertThat(atom.getDownloadStatus()).isEqualTo(DownloadStatus.SUCCEEDED); assertThat(atom.getModelType()).isEqualTo(MODEL_TYPE_ATOM); assertThat(atom.getUrlSuffix()).isEqualTo(MANIFEST_URL); assertThat(atom.getRunAttemptCount()).isEqualTo(RUN_ATTEMPT_COUNT); assertThat(atom.getFailureReason()).isEqualTo(FailureReason.UNKNOWN_FAILURE_REASON); assertThat(atom.getDownloadDurationMillis()).isEqualTo(DOWNLOAD_STARTED_TO_ENDED_MILLIS); } private void verifyFailedDownloadLogging() throws Exception { TextClassifierDownloadReported atom = Iterables.getOnlyElement(loggerTestRule.getLoggedDownloadReportedAtoms()); assertThat(atom.getWorkId()).isEqualTo(WORK_ID); assertThat(atom.getDownloadStatus()).isEqualTo(DownloadStatus.FAILED_AND_RETRY); assertThat(atom.getModelType()).isEqualTo(MODEL_TYPE_ATOM); assertThat(atom.getUrlSuffix()).isEqualTo(MANIFEST_URL); assertThat(atom.getRunAttemptCount()).isEqualTo(RUN_ATTEMPT_COUNT); assertThat(atom.getFailureReason()).isEqualTo(FAILED_TO_DOWNLOAD_FAILURE_REASON); assertThat(atom.getDownloaderLibFailureCode()) .isEqualTo(ModelDownloadException.DEFAULT_DOWNLOADER_LIB_ERROR_CODE); assertThat(atom.getDownloadDurationMillis()).isEqualTo(DOWNLOAD_STARTED_TO_ENDED_MILLIS); } private void verifyWorkLogging(int runTimeAttempt, WorkResult workResult) throws Exception { TextClassifierDownloadWorkCompleted atom = Iterables.getOnlyElement(loggerTestRule.getLoggedDownloadWorkCompletedAtoms()); assertThat(atom.getWorkId()).isEqualTo(WORK_ID); assertThat(atom.getWorkResult()).isEqualTo(workResult); assertThat(atom.getRunAttemptCount()).isEqualTo(runTimeAttempt); assertThat(atom.getWorkScheduledToStartedDurationMillis()) .isEqualTo(WORK_SCHEDULED_TO_STARTED_MILLIS); assertThat(atom.getWorkStartedToEndedDurationMillis()).isEqualTo(WORK_STARTED_TO_ENDED_MILLIS); } private void setUpManifestUrl( @ModelType.ModelTypeDef String modelType, String localeTag, String url) { String deviceConfigFlag = String.format(TextClassifierSettings.MANIFEST_URL_TEMPLATE, modelType, localeTag); deviceConfig.setConfig(deviceConfigFlag, url); } }