work/workmanager/src/androidTest/java/androidx/work/impl/WorkerWrapperTest.java - platform/frameworks/support - Git at Google

 /*
  * Copyright 2017 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 androidx.work.impl;

 import static androidx.work.WorkInfo.State.BLOCKED;
 import static androidx.work.WorkInfo.State.CANCELLED;
 import static androidx.work.WorkInfo.State.ENQUEUED;
 import static androidx.work.WorkInfo.State.FAILED;
 import static androidx.work.WorkInfo.State.RUNNING;
 import static androidx.work.WorkInfo.State.SUCCEEDED;

 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.notNullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.contains;
 import static org.hamcrest.Matchers.containsInAnyOrder;
 import static org.hamcrest.Matchers.greaterThan;
 import static org.hamcrest.Matchers.isOneOf;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.atLeast;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;

 import android.content.Context;
 import android.net.Uri;
 import android.os.Build;
 import android.util.Log;

 import androidx.test.core.app.ApplicationProvider;
 import androidx.test.ext.junit.runners.AndroidJUnit4;
 import androidx.test.filters.LargeTest;
 import androidx.test.filters.SmallTest;
 import androidx.work.ArrayCreatingInputMerger;
 import androidx.work.Configuration;
 import androidx.work.Data;
 import androidx.work.DatabaseTest;
 import androidx.work.ListenableWorker;
 import androidx.work.OneTimeWorkRequest;
 import androidx.work.PeriodicWorkRequest;
 import androidx.work.ProgressUpdater;
 import androidx.work.WorkerParameters;
 import androidx.work.impl.model.Dependency;
 import androidx.work.impl.model.DependencyDao;
 import androidx.work.impl.model.WorkSpec;
 import androidx.work.impl.model.WorkSpecDao;
 import androidx.work.impl.utils.SynchronousExecutor;
 import androidx.work.impl.utils.taskexecutor.InstantWorkTaskExecutor;
 import androidx.work.impl.utils.taskexecutor.TaskExecutor;
 import androidx.work.worker.ChainedArgumentWorker;
 import androidx.work.worker.EchoingWorker;
 import androidx.work.worker.ExceptionWorker;
 import androidx.work.worker.FailureWorker;
 import androidx.work.worker.InfiniteTestWorker;
 import androidx.work.worker.InterruptionAwareWorker;
 import androidx.work.worker.LatchWorker;
 import androidx.work.worker.RetryWorker;
 import androidx.work.worker.ReturnNullResultWorker;
 import androidx.work.worker.SleepTestWorker;
 import androidx.work.worker.TestWorker;
 import androidx.work.worker.UsedWorker;

 import com.google.common.util.concurrent.ListenableFuture;

 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.mockito.ArgumentCaptor;

 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executor;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;

 @RunWith(AndroidJUnit4.class)
 public class WorkerWrapperTest extends DatabaseTest {

  private Configuration mConfiguration;
  private TaskExecutor mWorkTaskExecutor;
  private WorkSpecDao mWorkSpecDao;
  private DependencyDao mDependencyDao;
  private Context mContext;
  private Scheduler mMockScheduler;
  private ProgressUpdater mMockProgressUpdater;
  private Executor mSynchronousExecutor = new SynchronousExecutor();

  @Before
  public void setUp() {
  mContext = ApplicationProvider.getApplicationContext();
  mConfiguration = new Configuration.Builder()
  .setExecutor(new SynchronousExecutor())
  .setMinimumLoggingLevel(Log.VERBOSE)
  .build();
  mWorkTaskExecutor = new InstantWorkTaskExecutor();
  mWorkSpecDao = mDatabase.workSpecDao();
  mDependencyDao = mDatabase.dependencyDao();
  mMockScheduler = mock(Scheduler.class);
  mMockProgressUpdater = mock(ProgressUpdater.class);
  }

  @Test
  @SmallTest
  public void testSuccess() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(SUCCEEDED));
  }

  @Test
  @SmallTest
  public void testRunAttemptCountIncremented_successfulExecution() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  insertWork(work);
  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build()
  .run();
  WorkSpec latestWorkSpec = mWorkSpecDao.getWorkSpec(work.getStringId());
  assertThat(latestWorkSpec.runAttemptCount, is(1));
  }

  @Test
  @SmallTest
  public void testRunAttemptCountIncremented_failedExecution() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(FailureWorker.class).build();
  insertWork(work);
  createBuilder(work.getStringId())
  .build()
  .run();
  WorkSpec latestWorkSpec = mWorkSpecDao.getWorkSpec(work.getStringId());
  assertThat(latestWorkSpec.runAttemptCount, is(1));
  }

  @Test
  @SmallTest
  public void testPermanentErrorWithInvalidWorkSpecId() {
  final String invalidWorkSpecId = "INVALID_ID";
  WorkerWrapper workerWrapper = createBuilder(invalidWorkSpecId).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  }

  @Test
  @SmallTest
  public void testInvalidWorkerClassName() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  work.getWorkSpec().workerClassName = "dummy";
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  verify(mMockScheduler, never()).schedule(any(WorkSpec[].class));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  @Test
  @SmallTest
  public void testUsedWorker_failsExecution() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  insertWork(work);

  UsedWorker usedWorker = (UsedWorker) mConfiguration.getWorkerFactory()
  .createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  UsedWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  Data.EMPTY,
  work.getTags(),
  new WorkerParameters.RuntimeExtras(),
  1,
  mSynchronousExecutor,
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));

  WorkerWrapper workerWrapper = createBuilder(work.getStringId())
  .withWorker(usedWorker)
  .build();
  workerWrapper.run();
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  @Test
  @SmallTest
  public void testNotEnqueued() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(RUNNING)
  .build();
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(true));
  }

  @Test
  @SmallTest
  public void testCancelled() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(CANCELLED)
  .build();
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(CANCELLED));
  }

  @Test
  @SmallTest
  public void testPermanentErrorWithInvalidWorkerClass() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  getWorkSpec(work).workerClassName = "INVALID_CLASS_NAME";
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  @Test
  @SmallTest
  public void testPermanentErrorWithInvalidInputMergerClass() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  getWorkSpec(work).inputMergerClassName = "INVALID_CLASS_NAME";
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  @Test
  @SmallTest
  public void testFailed() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(FailureWorker.class).build();
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  @Test
  @LargeTest
  public void testFailedOnDeepHierarchy() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(FailureWorker.class).build();
  insertWork(work);
  String previousId = work.getStringId();
  String firstWorkId = previousId;
  for (int i = 0; i < 500; ++i) {
  work = new OneTimeWorkRequest.Builder(FailureWorker.class).build();
  insertWork(work);
  mDependencyDao.insertDependency(new Dependency(work.getStringId(), previousId));
  previousId = work.getStringId();
  }
  WorkerWrapper workerWrapper = createBuilder(firstWorkId).build();
  workerWrapper.setFailedAndResolve();
  assertThat(mWorkSpecDao.getState(firstWorkId), is(FAILED));
  assertThat(mWorkSpecDao.getState(previousId), is(FAILED));
  }

  @Test
  @LargeTest
  public void testRunning() throws InterruptedException {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(SleepTestWorker.class).build();
  insertWork(work);
  WorkerWrapper wrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(wrapper);
  Executors.newSingleThreadExecutor().submit(wrapper);
  Thread.sleep(2000L); // Async wait duration.
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(RUNNING));
  Thread.sleep(SleepTestWorker.SLEEP_DURATION);
  assertThat(listener.mResult, is(false));
  }

  @Test
  @SmallTest
  public void testRunning_onlyWhenEnqueued() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(RUNNING)
  .build();
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(true));
  }

  @Test
  @SmallTest
  public void testDependencies() {
  OneTimeWorkRequest prerequisiteWork =
  new OneTimeWorkRequest.Builder(TestWorker.class).build();
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(BLOCKED).build();
  Dependency dependency = new Dependency(work.getStringId(), prerequisiteWork.getStringId());

  mDatabase.beginTransaction();
  try {
  insertWork(prerequisiteWork);
  insertWork(work);
  mDependencyDao.insertDependency(dependency);
  mDatabase.setTransactionSuccessful();
  } finally {
  mDatabase.endTransaction();
  }

  assertThat(mWorkSpecDao.getState(prerequisiteWork.getStringId()), is(ENQUEUED));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(BLOCKED));
  assertThat(mDependencyDao.hasCompletedAllPrerequisites(work.getStringId()), is(false));

  createBuilder(prerequisiteWork.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build()
  .run();

  assertThat(mWorkSpecDao.getState(prerequisiteWork.getStringId()), is(SUCCEEDED));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(ENQUEUED));
  assertThat(mDependencyDao.hasCompletedAllPrerequisites(work.getStringId()), is(true));

  ArgumentCaptor<WorkSpec> captor = ArgumentCaptor.forClass(WorkSpec.class);
  verify(mMockScheduler).schedule(captor.capture());
  assertThat(captor.getValue().id, is(work.getStringId()));
  }

  @Test
  @SmallTest
  public void testDependencies_passesOutputs() {
  OneTimeWorkRequest prerequisiteWork =
  new OneTimeWorkRequest.Builder(ChainedArgumentWorker.class).build();
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(BLOCKED)
  .build();
  Dependency dependency = new Dependency(work.getStringId(), prerequisiteWork.getStringId());

  mDatabase.beginTransaction();
  try {
  insertWork(prerequisiteWork);
  insertWork(work);
  mDependencyDao.insertDependency(dependency);
  mDatabase.setTransactionSuccessful();
  } finally {
  mDatabase.endTransaction();
  }

  createBuilder(prerequisiteWork.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build().run();

  List<Data> arguments = mWorkSpecDao.getInputsFromPrerequisites(work.getStringId());
  assertThat(arguments.size(), is(1));
  assertThat(arguments, contains(ChainedArgumentWorker.getChainedArguments()));
  }

  @Test
  @SmallTest
  public void testDependencies_passesMergedOutputs() {
  String key = "key";
  String value1 = "value1";
  String value2 = "value2";

  OneTimeWorkRequest prerequisiteWork1 = new OneTimeWorkRequest.Builder(EchoingWorker.class)
  .setInputData(new Data.Builder().putString(key, value1).build())
  .build();
  OneTimeWorkRequest prerequisiteWork2 = new OneTimeWorkRequest.Builder(EchoingWorker.class)
  .setInputData(new Data.Builder().putString(key, value2).build())
  .build();
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInputMerger(ArrayCreatingInputMerger.class)
  .build();
  Dependency dependency1 =
  new Dependency(work.getStringId(), prerequisiteWork1.getStringId());
  Dependency dependency2 =
  new Dependency(work.getStringId(), prerequisiteWork2.getStringId());

  mDatabase.beginTransaction();
  try {
  insertWork(prerequisiteWork1);
  insertWork(prerequisiteWork2);
  insertWork(work);
  mDependencyDao.insertDependency(dependency1);
  mDependencyDao.insertDependency(dependency2);
  mDatabase.setTransactionSuccessful();
  } finally {
  mDatabase.endTransaction();
  }

  // Run the prerequisites.
  createBuilder(prerequisiteWork1.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build().run();

  createBuilder(prerequisiteWork2.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build().run();

  // Create and run the dependent work.
  WorkerWrapper workerWrapper =
  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();
  workerWrapper.run();

  Data input = workerWrapper.mWorker.getInputData();
  assertThat(input.size(), is(1));
  assertThat(Arrays.asList(input.getStringArray(key)),
  containsInAnyOrder(value1, value2));
  }

  @Test
  @SmallTest
  public void testDependencies_setsPeriodStartTimesForUnblockedWork() {
  OneTimeWorkRequest prerequisiteWork =
  new OneTimeWorkRequest.Builder(TestWorker.class).build();
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(BLOCKED)
  .build();
  Dependency dependency = new Dependency(work.getStringId(), prerequisiteWork.getStringId());

  mDatabase.beginTransaction();
  try {
  insertWork(prerequisiteWork);
  insertWork(work);
  mDependencyDao.insertDependency(dependency);
  mDatabase.setTransactionSuccessful();
  } finally {
  mDatabase.endTransaction();
  }

  long beforeUnblockedTime = System.currentTimeMillis();

  createBuilder(prerequisiteWork.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build()
  .run();

  WorkSpec workSpec = mWorkSpecDao.getWorkSpec(work.getStringId());
  assertThat(workSpec.periodStartTime, is(greaterThan(beforeUnblockedTime)));
  }

  @Test
  @SmallTest
  public void testDependencies_enqueuesBlockedDependentsOnSuccess() {
  OneTimeWorkRequest prerequisiteWork =
  new OneTimeWorkRequest.Builder(TestWorker.class).build();
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(BLOCKED)
  .build();
  OneTimeWorkRequest cancelledWork = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(CANCELLED)
  .build();
  Dependency dependency1 = new Dependency(work.getStringId(), prerequisiteWork.getStringId());
  Dependency dependency2 =
  new Dependency(cancelledWork.getStringId(), prerequisiteWork.getStringId());

  mDatabase.beginTransaction();
  try {
  insertWork(prerequisiteWork);
  insertWork(work);
  insertWork(cancelledWork);
  mDependencyDao.insertDependency(dependency1);
  mDependencyDao.insertDependency(dependency2);
  mDatabase.setTransactionSuccessful();
  } finally {
  mDatabase.endTransaction();
  }

  createBuilder(prerequisiteWork.getStringId())
  .build()
  .run();

  assertThat(mWorkSpecDao.getState(prerequisiteWork.getStringId()), is(SUCCEEDED));
  assertThat(mWorkSpecDao.getState(work.getStringId()),
  isOneOf(ENQUEUED, RUNNING, SUCCEEDED));
  assertThat(mWorkSpecDao.getState(cancelledWork.getStringId()), is(CANCELLED));
  }

  @Test
  @SmallTest
  public void testDependencies_failsUncancelledDependentsOnFailure() {
  OneTimeWorkRequest prerequisiteWork =
  new OneTimeWorkRequest.Builder(FailureWorker.class).build();
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(BLOCKED)
  .build();
  OneTimeWorkRequest cancelledWork = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(CANCELLED)
  .build();
  Dependency dependency1 = new Dependency(work.getStringId(), prerequisiteWork.getStringId());
  Dependency dependency2 =
  new Dependency(cancelledWork.getStringId(), prerequisiteWork.getStringId());

  mDatabase.beginTransaction();
  try {
  insertWork(prerequisiteWork);
  insertWork(work);
  insertWork(cancelledWork);
  mDependencyDao.insertDependency(dependency1);
  mDependencyDao.insertDependency(dependency2);
  mDatabase.setTransactionSuccessful();
  } finally {
  mDatabase.endTransaction();
  }

  createBuilder(prerequisiteWork.getStringId())
  .build()
  .run();

  assertThat(mWorkSpecDao.getState(prerequisiteWork.getStringId()), is(FAILED));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  assertThat(mWorkSpecDao.getState(cancelledWork.getStringId()), is(CANCELLED));
  }

  @Test
  @SmallTest
  public void testBackedOffOneTimeWork_doesNotRun() {
  OneTimeWorkRequest retryWork =
  new OneTimeWorkRequest.Builder(RetryWorker.class).build();

  long future = System.currentTimeMillis() + TimeUnit.HOURS.toMillis(1);
  mDatabase.beginTransaction();
  try {
  mWorkSpecDao.insertWorkSpec(retryWork.getWorkSpec());
  mWorkSpecDao.setPeriodStartTime(retryWork.getStringId(), future);
  mWorkSpecDao.incrementWorkSpecRunAttemptCount(retryWork.getStringId());
  mDatabase.setTransactionSuccessful();
  } finally {
  mDatabase.endTransaction();
  }

  createBuilder(retryWork.getStringId())
  .build()
  .run();

  WorkSpec workSpec = mWorkSpecDao.getWorkSpec(retryWork.getStringId());
  // The run attempt count should remain the same
  assertThat(workSpec.runAttemptCount, is(1));
  }

  @Test
  @SmallTest
  public void testRun_periodicWork_success_updatesPeriodStartTime() {
  long intervalDuration = PeriodicWorkRequest.MIN_PERIODIC_INTERVAL_MILLIS;
  long periodStartTime = System.currentTimeMillis();

  PeriodicWorkRequest periodicWork = new PeriodicWorkRequest.Builder(
  TestWorker.class, intervalDuration, TimeUnit.MILLISECONDS).build();

  getWorkSpec(periodicWork).periodStartTime = periodStartTime;
  insertWork(periodicWork);

  createBuilder(periodicWork.getStringId())
  .build()
  .run();

  WorkSpec updatedWorkSpec = mWorkSpecDao.getWorkSpec(periodicWork.getStringId());
  assertThat(updatedWorkSpec.calculateNextRunTime(), greaterThan(periodStartTime));
  }

  @Test
  @SmallTest
  public void testRun_periodicWork_failure_updatesPeriodStartTime() {
  long intervalDuration = PeriodicWorkRequest.MIN_PERIODIC_INTERVAL_MILLIS;
  long periodStartTime = System.currentTimeMillis();

  PeriodicWorkRequest periodicWork = new PeriodicWorkRequest.Builder(
  FailureWorker.class, intervalDuration, TimeUnit.MILLISECONDS).build();

  getWorkSpec(periodicWork).periodStartTime = periodStartTime;
  insertWork(periodicWork);

  createBuilder(periodicWork.getStringId())
  .build()
  .run();

  WorkSpec updatedWorkSpec = mWorkSpecDao.getWorkSpec(periodicWork.getStringId());
  assertThat(updatedWorkSpec.calculateNextRunTime(), greaterThan(periodStartTime));
  }

  @Test
  @SmallTest
  public void testPeriodicWork_success() {
  PeriodicWorkRequest periodicWork = new PeriodicWorkRequest.Builder(
  TestWorker.class,
  PeriodicWorkRequest.MIN_PERIODIC_INTERVAL_MILLIS,
  TimeUnit.MILLISECONDS)
  .build();

  final String periodicWorkId = periodicWork.getStringId();
  insertWork(periodicWork);
  WorkerWrapper workerWrapper = createBuilder(periodicWorkId)
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();

  WorkSpec periodicWorkSpecAfterFirstRun = mWorkSpecDao.getWorkSpec(periodicWorkId);
  assertThat(listener.mResult, is(false));
  assertThat(periodicWorkSpecAfterFirstRun.runAttemptCount, is(0));
  assertThat(periodicWorkSpecAfterFirstRun.state, is(ENQUEUED));
  // SystemAlarmScheduler needs to reschedule the same worker.
  if (Build.VERSION.SDK_INT <= WorkManagerImpl.MAX_PRE_JOB_SCHEDULER_API_LEVEL) {
  ArgumentCaptor<WorkSpec> captor = ArgumentCaptor.forClass(WorkSpec.class);
  verify(mMockScheduler, atLeast(1))
  .schedule(captor.capture());

  WorkSpec workSpec = captor.getValue();
  assertThat(workSpec.id, is(periodicWorkId));
  }
  }

  @Test
  @SmallTest
  public void testPeriodicWork_fail() {
  PeriodicWorkRequest periodicWork = new PeriodicWorkRequest.Builder(
  FailureWorker.class,
  PeriodicWorkRequest.MIN_PERIODIC_INTERVAL_MILLIS,
  TimeUnit.MILLISECONDS)
  .build();

  final String periodicWorkId = periodicWork.getStringId();
  insertWork(periodicWork);
  WorkerWrapper workerWrapper = createBuilder(periodicWorkId).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();

  WorkSpec periodicWorkSpecAfterFirstRun = mWorkSpecDao.getWorkSpec(periodicWorkId);
  assertThat(listener.mResult, is(false));
  assertThat(periodicWorkSpecAfterFirstRun.runAttemptCount, is(0));
  assertThat(periodicWorkSpecAfterFirstRun.state, is(ENQUEUED));
  }

  @Test
  @SmallTest
  public void testPeriodicWork_retry() {
  PeriodicWorkRequest periodicWork = new PeriodicWorkRequest.Builder(
  RetryWorker.class,
  PeriodicWorkRequest.MIN_PERIODIC_INTERVAL_MILLIS,
  TimeUnit.MILLISECONDS)
  .build();

  final String periodicWorkId = periodicWork.getStringId();
  insertWork(periodicWork);
  WorkerWrapper workerWrapper = createBuilder(periodicWorkId).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();

  WorkSpec periodicWorkSpecAfterFirstRun = mWorkSpecDao.getWorkSpec(periodicWorkId);
  assertThat(listener.mResult, is(true));
  assertThat(periodicWorkSpecAfterFirstRun.runAttemptCount, is(1));
  assertThat(periodicWorkSpecAfterFirstRun.state, is(ENQUEUED));
  }


  @Test
  @SmallTest
  public void testPeriodic_dedupe() {
  PeriodicWorkRequest periodicWork = new PeriodicWorkRequest.Builder(
  TestWorker.class,
  PeriodicWorkRequest.MIN_PERIODIC_INTERVAL_MILLIS,
  TimeUnit.MILLISECONDS)
  .build();

  final String periodicWorkId = periodicWork.getStringId();
  final WorkSpec workSpec = periodicWork.getWorkSpec();
  long now = System.currentTimeMillis();
  workSpec.periodStartTime = now + workSpec.intervalDuration;
  insertWork(periodicWork);
  WorkerWrapper workerWrapper = createBuilder(periodicWorkId).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  // Should get rescheduled
  assertThat(listener.mResult, is(true));
  }

  @Test
  @SmallTest
  public void testPeriodic_firstRun_flexApplied_noDedupe() {
  PeriodicWorkRequest periodicWork = new PeriodicWorkRequest.Builder(
  TestWorker.class,
  PeriodicWorkRequest.MIN_PERIODIC_INTERVAL_MILLIS,
  TimeUnit.MILLISECONDS,
  PeriodicWorkRequest.MIN_PERIODIC_FLEX_MILLIS,
  TimeUnit.MILLISECONDS)
  .build();

  final String periodicWorkId = periodicWork.getStringId();
  final WorkSpec workSpec = periodicWork.getWorkSpec();
  workSpec.periodStartTime = 0;
  insertWork(periodicWork);
  WorkerWrapper workerWrapper = createBuilder(periodicWorkId).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  // Should not get rescheduled
  assertThat(listener.mResult, is(false));
  }

  @Test
  @SmallTest
  public void testScheduler() {
  OneTimeWorkRequest prerequisiteWork =
  new OneTimeWorkRequest.Builder(TestWorker.class).build();
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class)
  .setInitialState(BLOCKED).build();
  Dependency dependency = new Dependency(work.getStringId(), prerequisiteWork.getStringId());

  mDatabase.beginTransaction();
  try {
  insertWork(prerequisiteWork);
  insertWork(work);
  mDependencyDao.insertDependency(dependency);
  mDatabase.setTransactionSuccessful();
  } finally {
  mDatabase.endTransaction();
  }

  createBuilder(prerequisiteWork.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build()
  .run();

  ArgumentCaptor<WorkSpec> captor = ArgumentCaptor.forClass(WorkSpec.class);
  verify(mMockScheduler).schedule(captor.capture());
  assertThat(captor.getValue().id, is(work.getStringId()));
  }

  @Test
  @SmallTest
  public void testFromWorkSpec_hasAppContext() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  ListenableWorker worker = mConfiguration.getWorkerFactory().createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  TestWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  Data.EMPTY,
  work.getTags(),
  new WorkerParameters.RuntimeExtras(),
  1,
  mSynchronousExecutor,
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));

  assertThat(worker, is(notNullValue()));
  assertThat(worker.getApplicationContext(), is(equalTo(mContext.getApplicationContext())));
  }

  @Test
  @SmallTest
  public void testFromWorkSpec_hasCorrectArguments() {
  String key = "KEY";
  String expectedValue = "VALUE";
  Data input = new Data.Builder().putString(key, expectedValue).build();

  OneTimeWorkRequest work =
  new OneTimeWorkRequest.Builder(TestWorker.class).setInputData(input).build();
  ListenableWorker worker = mConfiguration.getWorkerFactory().createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  TestWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  input,
  work.getTags(),
  new WorkerParameters.RuntimeExtras(),
  1,
  mSynchronousExecutor,
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));

  assertThat(worker, is(notNullValue()));
  assertThat(worker.getInputData().getString(key), is(expectedValue));

  work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  worker = mConfiguration.getWorkerFactory().createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  TestWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  Data.EMPTY,
  work.getTags(),
  new WorkerParameters.RuntimeExtras(),
  1,
  mSynchronousExecutor,
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));

  assertThat(worker, is(notNullValue()));
  assertThat(worker.getInputData().size(), is(0));
  }

  @Test
  @SmallTest
  public void testFromWorkSpec_hasCorrectTags() {
  OneTimeWorkRequest work =
  new OneTimeWorkRequest.Builder(TestWorker.class)
  .addTag("one")
  .addTag("two")
  .addTag("three")
  .build();
  ListenableWorker worker = mConfiguration.getWorkerFactory().createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  TestWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  Data.EMPTY,
  Arrays.asList("one", "two", "three"),
  new WorkerParameters.RuntimeExtras(),
  1,
  mSynchronousExecutor,
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));

  assertThat(worker, is(notNullValue()));
  assertThat(worker.getTags(), containsInAnyOrder("one", "two", "three"));
  }

  @Test
  @SmallTest
  public void testFromWorkSpec_hasCorrectRuntimeExtras() {
  OneTimeWorkRequest work =
  new OneTimeWorkRequest.Builder(TestWorker.class).build();
  WorkerParameters.RuntimeExtras runtimeExtras = new WorkerParameters.RuntimeExtras();
  runtimeExtras.triggeredContentAuthorities = Arrays.asList("tca1", "tca2", "tca3");
  runtimeExtras.triggeredContentUris = Arrays.asList(Uri.parse("tcu1"), Uri.parse("tcu2"));

  ListenableWorker worker = mConfiguration.getWorkerFactory().createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  TestWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  Data.EMPTY,
  work.getTags(),
  runtimeExtras,
  1,
  mSynchronousExecutor,
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));

  assertThat(worker, is(notNullValue()));
  assertThat(worker.getTriggeredContentAuthorities(),
  containsInAnyOrder(runtimeExtras.triggeredContentAuthorities.toArray()));
  assertThat(worker.getTriggeredContentUris(),
  containsInAnyOrder(runtimeExtras.triggeredContentUris.toArray()));
  }

  @Test
  @SmallTest
  public void testSuccess_withPendingScheduledWork() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  insertWork(work);

  OneTimeWorkRequest unscheduled = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  insertWork(unscheduled);

  WorkerWrapper workerWrapper = createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();

  verify(mMockScheduler, times(1)).schedule(unscheduled.getWorkSpec());
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(SUCCEEDED));
  }

  @Test
  @SmallTest
  public void testFailure_withPendingScheduledWork() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(FailureWorker.class).build();
  insertWork(work);

  OneTimeWorkRequest unscheduled = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  insertWork(unscheduled);

  WorkerWrapper workerWrapper = createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();

  verify(mMockScheduler, times(1)).schedule(unscheduled.getWorkSpec());
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  @Test
  @LargeTest
  public void testInterruption() throws InterruptedException {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(TestWorker.class).build();
  insertWork(work);

  WorkerWrapper workerWrapper =
  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  Executors.newSingleThreadExecutor().submit(workerWrapper);
  workerWrapper.interrupt(false);
  Thread.sleep(6000L);
  assertThat(listener.mResult, is(true));
  }

  @Test
  @LargeTest
  public void testPruneWhileRunning_callsSchedulerCancel() throws InterruptedException {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(LatchWorker.class).build();
  insertWork(work);

  LatchWorker latchWorker =
  (LatchWorker) mConfiguration.getWorkerFactory().createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  LatchWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  Data.EMPTY,
  work.getTags(),
  new WorkerParameters.RuntimeExtras(),
  1,
  Executors.newSingleThreadExecutor(),
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));

  WorkerWrapper workerWrapper =
  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .withWorker(latchWorker)
  .build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  Executors.newSingleThreadExecutor().submit(workerWrapper);

  Thread.sleep(1000L);

  mDatabase.workSpecDao().delete(work.getStringId());
  assertThat(latchWorker.mLatch.getCount(), is(greaterThan(0L)));

  latchWorker.mLatch.countDown();

  Thread.sleep(1000L);

  assertThat(listener.mResult, is(notNullValue()));
  verify(mMockScheduler, times(1)).cancel(work.getStringId());
  }

  @Test
  @SmallTest
  public void testInterruptionWithoutCancellation_isMarkedOnRunningWorker() {
  OneTimeWorkRequest work =
  new OneTimeWorkRequest.Builder(InterruptionAwareWorker.class).build();
  insertWork(work);

  ListenableWorker worker = mConfiguration.getWorkerFactory().createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  InterruptionAwareWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  Data.EMPTY,
  work.getTags(),
  new WorkerParameters.RuntimeExtras(),
  1,
  mSynchronousExecutor,
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));
  assertThat(worker, is(notNullValue()));
  assertThat(worker.isStopped(), is(false));

  WorkerWrapper workerWrapper =
  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .withWorker(worker)
  .build();
  Executors.newSingleThreadExecutor().submit(workerWrapper);
  workerWrapper.interrupt(false);
  assertThat(worker.isStopped(), is(true));
  }

  @Test
  @SmallTest
  public void testInterruptionWithCancellation_isMarkedOnRunningWorker() {
  OneTimeWorkRequest work =
  new OneTimeWorkRequest.Builder(InterruptionAwareWorker.class).build();
  insertWork(work);

  ListenableWorker worker = mConfiguration.getWorkerFactory().createWorkerWithDefaultFallback(
  mContext.getApplicationContext(),
  InterruptionAwareWorker.class.getName(),
  new WorkerParameters(
  work.getId(),
  Data.EMPTY,
  Collections.<String>emptyList(),
  new WorkerParameters.RuntimeExtras(),
  1,
  mSynchronousExecutor,
  mWorkTaskExecutor,
  mConfiguration.getWorkerFactory(),
  mMockProgressUpdater));
  assertThat(worker, is(notNullValue()));
  assertThat(worker.isStopped(), is(false));

  WorkerWrapper workerWrapper =
  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .withWorker(worker)
  .build();
  Executors.newSingleThreadExecutor().submit(workerWrapper);
  workerWrapper.interrupt(true);
  assertThat(worker.isStopped(), is(true));
  }

  @Test
  @SmallTest
  public void testException_isTreatedAsFailure() {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(ExceptionWorker.class).build();
  insertWork(work);

  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build()
  .run();

  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  @Test
  @LargeTest
  public void testWorkerWrapper_handlesWorkSpecDeletion() throws InterruptedException {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(SleepTestWorker.class).build();
  insertWork(work);

  WorkerWrapper workerWrapper =
  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  Executors.newSingleThreadExecutor().submit(workerWrapper);
  mWorkSpecDao.delete(work.getStringId());
  Thread.sleep(6000L);
  assertThat(listener.mResult, is(false));
  }

  @Test
  @LargeTest
  public void testWorker_getsRunAttemptCount() throws InterruptedException {
  OneTimeWorkRequest work = new OneTimeWorkRequest.Builder(InfiniteTestWorker.class)
  .setInitialRunAttemptCount(10)
  .build();
  insertWork(work);

  WorkerWrapper workerWrapper =
  createBuilder(work.getStringId())
  .withSchedulers(Collections.singletonList(mMockScheduler))
  .build();

  Executors.newSingleThreadExecutor().submit(workerWrapper);
  Thread.sleep(1000L);
  assertThat(workerWrapper.mWorker.getRunAttemptCount(), is(10));
  }

  @Test
  @SmallTest
  public void testWorkerThatReturnsNullResult() {
  OneTimeWorkRequest work =
  new OneTimeWorkRequest.Builder(ReturnNullResultWorker.class).build();
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  @Test
  @SmallTest
  public void testWorkerThatThrowsAnException() {
  OneTimeWorkRequest work =
  new OneTimeWorkRequest.Builder(ExceptionWorker.class).build();
  insertWork(work);
  WorkerWrapper workerWrapper = createBuilder(work.getStringId()).build();
  FutureListener listener = createAndAddFutureListener(workerWrapper);
  workerWrapper.run();
  assertThat(listener.mResult, is(false));
  assertThat(mWorkSpecDao.getState(work.getStringId()), is(FAILED));
  }

  private WorkerWrapper.Builder createBuilder(String workSpecId) {
  return new WorkerWrapper.Builder(
  mContext,
  mConfiguration,
  mWorkTaskExecutor,
  mDatabase,
  workSpecId);
  }

  private FutureListener createAndAddFutureListener(WorkerWrapper workerWrapper) {
  ListenableFuture<Boolean> future = workerWrapper.getFuture();
  FutureListener listener = new FutureListener(future);
  future.addListener(listener, mSynchronousExecutor);
  return listener;
  }

  private static class FutureListener implements Runnable {

  ListenableFuture<Boolean> mFuture;
  Boolean mResult;

  FutureListener(ListenableFuture<Boolean> future) {
  mFuture = future;
  }

  @Override
  public void run() {
  try {
  mResult = mFuture.get();
  } catch (InterruptedException | ExecutionException e) {
  // Do nothing.
  }
  }
  }
 }