Keep track of completed work in TestScheduler.

* This is a workaround for incorrect behavior for tests where we did
correctly switch to a synchronous internal task executor thereby
requiring that users call setInitialDelayMet() or setAllConstraintsMet().

* Now the TestScheduler tracks worker completion and does not throw
an exception for WorkSpecs that it knows about.

Test: Added unit tests.
Fixes: b/139190782
Change-Id: I42b53338650c774ba26fdbeffa5c037067786cf2

diff --git a/work/workmanager-testing/src/androidTest/java/androidx/work/testing/TestSchedulerTest.java b/work/workmanager-testing/src/androidTest/java/androidx/work/testing/TestSchedulerTest.java
index 3bbb0bb..868ce63 100644
--- a/work/workmanager-testing/src/androidTest/java/androidx/work/testing/TestSchedulerTest.java
+++ b/work/workmanager-testing/src/androidTest/java/androidx/work/testing/TestSchedulerTest.java

@@ -161,6 +161,31 @@
         }
     }
 
+    @Test
+    public void testWorker_afterSuccessfulRun_postConditions()
+            throws InterruptedException, ExecutionException {
+
+        OneTimeWorkRequest request = createWorkRequest();
+        WorkManager workManager = WorkManager.getInstance(mContext);
+        workManager.enqueue(request).getResult().get();
+        WorkInfo status = workManager.getWorkInfoById(request.getId()).get();
+        assertThat(status.getState().isFinished(), is(true));
+        mTestDriver.setAllConstraintsMet(request.getId());
+        mTestDriver.setInitialDelayMet(request.getId());
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testWorker_afterSuccessfulRun_throwsExceptionWhenSetPeriodDelayMet()
+            throws InterruptedException, ExecutionException {
+
+        OneTimeWorkRequest request = createWorkRequest();
+        WorkManager workManager = WorkManager.getInstance(mContext);
+        workManager.enqueue(request).getResult().get();
+        WorkInfo status = workManager.getWorkInfoById(request.getId()).get();
+        assertThat(status.getState().isFinished(), is(true));
+        mTestDriver.setPeriodDelayMet(request.getId());
+    }
+
     private static OneTimeWorkRequest createWorkRequest() {
         return new OneTimeWorkRequest.Builder(TestWorker.class).build();
     }

diff --git a/work/workmanager-testing/src/main/java/androidx/work/testing/TestScheduler.java b/work/workmanager-testing/src/main/java/androidx/work/testing/TestScheduler.java
index 4a526ce..afc4f41 100644
--- a/work/workmanager-testing/src/main/java/androidx/work/testing/TestScheduler.java
+++ b/work/workmanager-testing/src/main/java/androidx/work/testing/TestScheduler.java

@@ -47,13 +47,15 @@
 class TestScheduler implements Scheduler, ExecutionListener {
 
     private final Context mContext;
-    private final Map<String, InternalWorkState> mInternalWorkStates;
+    private final Map<String, InternalWorkState> mPendingWorkStates;
+    private final Map<String, InternalWorkState> mTerminatedWorkStates;
 
     private static final Object sLock = new Object();
 
     TestScheduler(@NonNull Context context) {
         mContext = context;
-        mInternalWorkStates = new HashMap<>();
+        mPendingWorkStates = new HashMap<>();
+        mTerminatedWorkStates = new HashMap<>();
     }
 
     @Override
@@ -65,8 +67,8 @@
         synchronized (sLock) {
             List<String> workSpecIdsToSchedule = new ArrayList<>(workSpecs.length);
             for (WorkSpec workSpec : workSpecs) {
-                if (!mInternalWorkStates.containsKey(workSpec.id)) {
-                    mInternalWorkStates.put(workSpec.id, new InternalWorkState(mContext, workSpec));
+                if (!mPendingWorkStates.containsKey(workSpec.id)) {
+                    mPendingWorkStates.put(workSpec.id, new InternalWorkState(mContext, workSpec));
                 }
                 workSpecIdsToSchedule.add(workSpec.id);
             }
@@ -78,7 +80,11 @@
     public void cancel(@NonNull String workSpecId) {
         synchronized (sLock) {
             WorkManagerImpl.getInstance(mContext).stopWork(workSpecId);
-            mInternalWorkStates.remove(workSpecId);
+            mPendingWorkStates.remove(workSpecId);
+            // We don't need to keep track of cancelled workSpecs. This is because subsequent calls
+            // to enqueue() will no-op because insertWorkSpec in WorkDatabase has a conflict
+            // policy of @Ignore. So TestScheduler will _never_ be asked to schedule those
+            // WorkSpecs.
         }
     }
 
@@ -91,13 +97,16 @@
      */
     void setAllConstraintsMet(@NonNull UUID workSpecId) {
         synchronized (sLock) {
-            InternalWorkState internalWorkState = mInternalWorkStates.get(workSpecId.toString());
-            if (internalWorkState == null) {
-                throw new IllegalArgumentException(
-                        "Work with id " + workSpecId + " is not enqueued!");
+            String id = workSpecId.toString();
+            if (!mTerminatedWorkStates.containsKey(id)) {
+                InternalWorkState internalWorkState = mPendingWorkStates.get(id);
+                if (internalWorkState == null) {
+                    throw new IllegalArgumentException(
+                            "Work with id " + workSpecId + " is not enqueued!");
+                }
+                internalWorkState.mConstraintsMet = true;
+                scheduleInternal(Collections.singletonList(workSpecId.toString()));
             }
-            internalWorkState.mConstraintsMet = true;
-            scheduleInternal(Collections.singletonList(workSpecId.toString()));
         }
     }
 
@@ -110,13 +119,16 @@
      */
     void setInitialDelayMet(@NonNull UUID workSpecId) {
         synchronized (sLock) {
-            InternalWorkState internalWorkState = mInternalWorkStates.get(workSpecId.toString());
-            if (internalWorkState == null) {
-                throw new IllegalArgumentException(
-                        "Work with id " + workSpecId + " is not enqueued!");
+            String id = workSpecId.toString();
+            if (!mTerminatedWorkStates.containsKey(id)) {
+                InternalWorkState internalWorkState = mPendingWorkStates.get(id);
+                if (internalWorkState == null) {
+                    throw new IllegalArgumentException(
+                            "Work with id " + workSpecId + " is not enqueued!");
+                }
+                internalWorkState.mInitialDelayMet = true;
+                scheduleInternal(Collections.singletonList(workSpecId.toString()));
             }
-            internalWorkState.mInitialDelayMet = true;
-            scheduleInternal(Collections.singletonList(workSpecId.toString()));
         }
     }
 
@@ -129,7 +141,8 @@
      */
     void setPeriodDelayMet(@NonNull UUID workSpecId) {
         synchronized (sLock) {
-            InternalWorkState internalWorkState = mInternalWorkStates.get(workSpecId.toString());
+            String id = workSpecId.toString();
+            InternalWorkState internalWorkState = mPendingWorkStates.get(id);
             if (internalWorkState == null) {
                 throw new IllegalArgumentException(
                         "Work with id " + workSpecId + " is not enqueued!");
@@ -141,20 +154,22 @@
 
     @Override
     public void onExecuted(@NonNull String workSpecId, boolean needsReschedule) {
-
         synchronized (sLock) {
-            InternalWorkState internalWorkState = mInternalWorkStates.get(workSpecId);
-            if (internalWorkState.mWorkSpec.isPeriodic()) {
-                internalWorkState.reset();
-            } else {
-                mInternalWorkStates.remove(workSpecId);
+            InternalWorkState internalWorkState = mPendingWorkStates.get(workSpecId);
+            if (internalWorkState != null) {
+                if (internalWorkState.mWorkSpec.isPeriodic()) {
+                    internalWorkState.reset();
+                } else {
+                    mTerminatedWorkStates.put(workSpecId, internalWorkState);
+                    mPendingWorkStates.remove(workSpecId);
+                }
             }
         }
     }
 
     private void scheduleInternal(Collection<String> workSpecIds) {
         for (String workSpecId : workSpecIds) {
-            InternalWorkState internalWorkState = mInternalWorkStates.get(workSpecId);
+            InternalWorkState internalWorkState = mPendingWorkStates.get(workSpecId);
             if (internalWorkState.isRunnable()) {
                 WorkManagerImpl.getInstance(mContext).startWork(workSpecId);
             }