Add new test for job timing out while work is enqueued.

Takes advantage of new job scheduler command to force an immediate
timeout.

Test: ran this.
Change-Id: I721dc8974ce3303ae94705db79464cdc701026e0
diff --git a/tests/JobScheduler/src/android/jobscheduler/MockJobService.java b/tests/JobScheduler/src/android/jobscheduler/MockJobService.java
index 1955483..4ade3b5 100644
--- a/tests/JobScheduler/src/android/jobscheduler/MockJobService.java
+++ b/tests/JobScheduler/src/android/jobscheduler/MockJobService.java
@@ -50,11 +50,14 @@
 
     private JobParameters mParams;
 
-    ArrayList<Intent> mReceivedWork = new ArrayList<Intent>();
+    ArrayList<JobWorkItem> mReceivedWork = new ArrayList<>();
+
+    private boolean mWaitingForStop;
 
     @Override
     public void onDestroy() {
         super.onDestroy();
+        Log.i(TAG, "Destroying test service");
         if (TestEnvironment.getTestEnvironment().getExpectedWork() != null) {
             TestEnvironment.getTestEnvironment().notifyExecution(mParams, 0, 0, mReceivedWork,
                     null);
@@ -99,7 +102,7 @@
             int index = 0;
             while ((work = params.dequeueWork()) != null) {
                 Log.i(TAG, "Received work #" + index + ": " + work.getIntent());
-                mReceivedWork.add(work.getIntent());
+                mReceivedWork.add(work);
 
                 if (index < expectedWork.length) {
                     TestWorkItem expected = expectedWork[index];
@@ -166,9 +169,16 @@
                             }
                         }
                     }
+
+                    if ((expected.flags & TestWorkItem.FLAG_WAIT_FOR_STOP) != 0) {
+                        Log.i(TAG, "Now waiting to stop");
+                        mWaitingForStop = true;
+                        TestEnvironment.getTestEnvironment().notifyWaitingForStop();
+                        return true;
+                    }
                 }
 
-                params.completeWork(work);
+                mParams.completeWork(work);
 
                 if (index < expectedWork.length) {
                     TestWorkItem expected = expectedWork[index];
@@ -198,12 +208,17 @@
 
     @Override
     public boolean onStopJob(JobParameters params) {
-        return false;
+        Log.i(TAG, "Received stop callback");
+        return mWaitingForStop;
     }
 
     public static final class TestWorkItem {
+        public static final int FLAG_WAIT_FOR_STOP = 1<<0;
+
         public final Intent intent;
         public final JobInfo jobInfo;
+        public final int flags;
+        public final int deliveryCount;
         public final TestWorkItem[] subitems;
         public final Uri[] requireUrisGranted;
         public final Uri[] requireUrisNotGranted;
@@ -211,6 +226,18 @@
         public TestWorkItem(Intent _intent) {
             intent = _intent;
             jobInfo = null;
+            flags = 0;
+            deliveryCount = 1;
+            subitems = null;
+            requireUrisGranted = null;
+            requireUrisNotGranted = null;
+        }
+
+        public TestWorkItem(Intent _intent, int _flags, int _deliveryCount) {
+            intent = _intent;
+            jobInfo = null;
+            flags = _flags;
+            deliveryCount = _deliveryCount;
             subitems = null;
             requireUrisGranted = null;
             requireUrisNotGranted = null;
@@ -219,6 +246,8 @@
         public TestWorkItem(Intent _intent, JobInfo _jobInfo, TestWorkItem[] _subitems) {
             intent = _intent;
             jobInfo = _jobInfo;
+            flags = 0;
+            deliveryCount = 1;
             subitems = _subitems;
             requireUrisGranted = null;
             requireUrisNotGranted = null;
@@ -228,10 +257,17 @@
                 Uri[] _requireUrisNotGranted) {
             intent = _intent;
             jobInfo = null;
+            flags = 0;
+            deliveryCount = 1;
             subitems = null;
             requireUrisGranted = _requireUrisGranted;
             requireUrisNotGranted = _requireUrisNotGranted;
         }
+
+        @Override
+        public String toString() {
+            return "TestWorkItem { " + intent + " dc=" + deliveryCount + " }";
+        }
     }
 
     /**
@@ -245,12 +281,13 @@
         //public static final int INVALID_JOB_ID = -1;
 
         private CountDownLatch mLatch;
+        private CountDownLatch mWaitingForStopLatch;
         private CountDownLatch mDoWorkLatch;
         private TestWorkItem[] mExpectedWork;
         private JobParameters mExecutedJobParameters;
         private int mExecutedPermCheckRead;
         private int mExecutedPermCheckWrite;
-        private ArrayList<Intent> mExecutedReceivedWork;
+        private ArrayList<JobWorkItem> mExecutedReceivedWork;
         private String mExecutedErrorMessage;
 
         public static TestEnvironment getTestEnvironment() {
@@ -276,7 +313,7 @@
             return mExecutedPermCheckWrite;
         }
 
-        public ArrayList<Intent> getLastReceivedWork() {
+        public ArrayList<JobWorkItem> getLastReceivedWork() {
             return mExecutedReceivedWork;
         }
 
@@ -305,12 +342,16 @@
             return !mLatch.await(DEFAULT_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
         }
 
+        public boolean awaitWaitingForStop() throws InterruptedException {
+            return !mWaitingForStopLatch.await(DEFAULT_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
+        }
+
         public boolean awaitDoWork() throws InterruptedException {
             return !mDoWorkLatch.await(DEFAULT_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
         }
 
         private void notifyExecution(JobParameters params, int permCheckRead, int permCheckWrite,
-                ArrayList<Intent> receivedWork, String errorMsg) {
+                ArrayList<JobWorkItem> receivedWork, String errorMsg) {
             //Log.d(TAG, "Job executed:" + params.getJobId());
             mExecutedJobParameters = params;
             mExecutedPermCheckRead = permCheckRead;
@@ -320,6 +361,10 @@
             mLatch.countDown();
         }
 
+        private void notifyWaitingForStop() {
+            mWaitingForStopLatch.countDown();
+        }
+
         public void setExpectedExecutions(int numExecutions) {
             // For no executions expected, set count to 1 so we can still block for the timeout.
             if (numExecutions == 0) {
@@ -329,6 +374,10 @@
             }
         }
 
+        public void setExpectedWaitForStop() {
+            mWaitingForStopLatch = new CountDownLatch(1);
+        }
+
         public void setExpectedWork(TestWorkItem[] work) {
             mExpectedWork = work;
             mDoWorkLatch = new CountDownLatch(1);
diff --git a/tests/JobScheduler/src/android/jobscheduler/cts/EnqueueJobWorkTest.java b/tests/JobScheduler/src/android/jobscheduler/cts/EnqueueJobWorkTest.java
index 402233b..ff9a57b 100644
--- a/tests/JobScheduler/src/android/jobscheduler/cts/EnqueueJobWorkTest.java
+++ b/tests/JobScheduler/src/android/jobscheduler/cts/EnqueueJobWorkTest.java
@@ -24,6 +24,8 @@
 import android.jobscheduler.MockJobService.TestWorkItem;
 import android.net.Uri;
 
+import com.android.compatibility.common.util.SystemUtil;
+
 import java.util.ArrayList;
 
 /**
@@ -65,7 +67,7 @@
         return i1.filterEquals(i2);
     }
 
-    private void compareWork(TestWorkItem[] expected, ArrayList<Intent> received) {
+    private void compareWork(TestWorkItem[] expected, ArrayList<JobWorkItem> received) {
         if (received == null) {
             fail("Didn't receive any expected work.");
         }
@@ -74,7 +76,7 @@
             expectedArray.add(expected[i]);
         }
         for (int i = 0; i < received.size(); i++) {
-            Intent work = received.get(i);
+            JobWorkItem work = received.get(i);
             if (i < expected.length && expected[i].subitems != null) {
                 TestWorkItem[] sub = expected[i].subitems;
                 for (int j = 0; j < sub.length; j++) {
@@ -85,8 +87,13 @@
                 fail("Received more than " + expected.length + " work items, first extra is "
                         + work);
             }
-            if (!intentEquals(work, expectedArray.get(i).intent)) {
-                fail("Received work #" + i + " " + work + " but expected " + expected[i]);
+            if (!intentEquals(work.getIntent(), expectedArray.get(i).intent)) {
+                fail("Received work #" + i + " " + work.getIntent() + " but expected " + expected[i]);
+            }
+            if (work.getDeliveryCount() != expectedArray.get(i).deliveryCount) {
+                fail("Received work #" + i + " " + work.getIntent() + " delivery count is "
+                        + work.getDeliveryCount() + " but expected "
+                        + expectedArray.get(i).deliveryCount);
             }
         }
         if (received.size() < expected.length) {
@@ -179,6 +186,58 @@
     }
 
     /**
+     * Test job getting stopped while processing work and that work being redelivered.
+     */
+    public void testEnqueueMultipleRedeliver() throws Exception {
+        Intent work1 = new Intent("work1");
+        Intent work2 = new Intent("work2");
+        Intent work3 = new Intent("work3");
+        Intent work4 = new Intent("work4");
+        Intent work5 = new Intent("work5");
+        Intent work6 = new Intent("work6");
+        TestWorkItem[] initialWork = new TestWorkItem[] {
+                new TestWorkItem(work1), new TestWorkItem(work2), new TestWorkItem(work3),
+                new TestWorkItem(work4, TestWorkItem.FLAG_WAIT_FOR_STOP, 1) };
+        kTestEnvironment.setExpectedExecutions(1);
+        kTestEnvironment.setExpectedWaitForStop();
+        kTestEnvironment.setExpectedWork(initialWork);
+        JobInfo ji = mBuilder.setOverrideDeadline(0).build();
+        mJobScheduler.enqueue(ji, new JobWorkItem(work1));
+        mJobScheduler.enqueue(ji, new JobWorkItem(work2));
+        mJobScheduler.enqueue(ji, new JobWorkItem(work3));
+        mJobScheduler.enqueue(ji, new JobWorkItem(work4));
+        kTestEnvironment.readyToWork();
+
+        // Now wait for the job to get to the point where it is processing the last
+        // work and waiting for it to be stopped.
+        assertFalse("Job with work enqueued did not wait to stop.",
+                kTestEnvironment.awaitWaitingForStop());
+
+        // Cause the job to timeout (stop) immediately, and wait for its execution to finish.
+        SystemUtil.runShellCommand(getInstrumentation(), "cmd jobscheduler timeout "
+                + kJobServiceComponent.getPackageName() + " " + ENQUEUE_WORK_JOB_ID);
+        assertTrue("Job with work enqueued did not finish.",
+                kTestEnvironment.awaitExecution());
+        compareWork(initialWork, kTestEnvironment.getLastReceivedWork());
+
+        // Now we are going to add some more work, restart the job, and see if it correctly
+        // redelivers the last work and delivers the new work.
+        TestWorkItem[] finalWork = new TestWorkItem[] {
+                new TestWorkItem(work4, 0, 2), new TestWorkItem(work5), new TestWorkItem(work6) };
+        kTestEnvironment.setExpectedExecutions(1);
+        kTestEnvironment.setExpectedWork(finalWork);
+        mJobScheduler.enqueue(ji, new JobWorkItem(work5));
+        mJobScheduler.enqueue(ji, new JobWorkItem(work6));
+        kTestEnvironment.readyToWork();
+        SystemUtil.runShellCommand(getInstrumentation(), "cmd jobscheduler run "
+                + kJobServiceComponent.getPackageName() + " " + ENQUEUE_WORK_JOB_ID);
+
+        assertTrue("Restarted with work enqueued did not execute.",
+                kTestEnvironment.awaitExecution());
+        compareWork(finalWork, kTestEnvironment.getLastReceivedWork());
+    }
+
+    /**
      * Test basic enqueueing batches of work.
      */
     public void testEnqueueMultipleUriGrantWork() throws Exception {