path: root/java/tests/instrumentation/src/com/android/textclassifier/downloader/ModelDownloadWorkerTest.java

/*
 * 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<TextClassifierDownloadReported> 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<TextClassifierDownloadReported> 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<TextClassifierDownloadReported> 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<TextClassifierDownloadReported> 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<TextClassifierDownloadReported> 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<TextClassifierDownloadReported> 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);
  }
}