Test cases for changed long-write behavior

It cannot happen anymore that we can press the print button while we
write (the PrintActivity does not have a print button during update), so
the original test case does not make sense anymore.

Add new test cases:
1 The print button should be gone while we update the preview
2 - The print button re-appears once an update works
  - We can change the printer during update

Test: cts-tradefed run cts-dev -m Print
Bug: 36599750
Change-Id: I6362b21f082869b96a33a78810f90d9b8bea7e53
diff --git a/tests/tests/print/src/android/print/cts/BasePrintTest.java b/tests/tests/print/src/android/print/cts/BasePrintTest.java
index 025a611..198fe9f 100644
--- a/tests/tests/print/src/android/print/cts/BasePrintTest.java
+++ b/tests/tests/print/src/android/print/cts/BasePrintTest.java
@@ -18,6 +18,7 @@
 
 import static android.print.cts.Utils.getPrintManager;
 
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
 import static org.junit.Assert.fail;
 import static org.junit.Assume.assumeTrue;
@@ -129,6 +130,7 @@
     private CallCounter mCancelOperationCounter;
     private CallCounter mLayoutCallCounter;
     private CallCounter mWriteCallCounter;
+    private CallCounter mWriteCancelCallCounter;
     private CallCounter mFinishCallCounter;
     private CallCounter mPrintJobQueuedCallCounter;
     private CallCounter mCreateSessionCallCounter;
@@ -244,6 +246,7 @@
         mLayoutCallCounter = new CallCounter();
         mFinishCallCounter = new CallCounter();
         mWriteCallCounter = new CallCounter();
+        mWriteCancelCallCounter = new CallCounter();
         mFinishCallCounter = new CallCounter();
         mPrintJobQueuedCallCounter = new CallCounter();
         mCreateSessionCallCounter = new CallCounter();
@@ -345,6 +348,10 @@
         mWriteCallCounter.call();
     }
 
+    protected void onWriteCancelCalled() {
+        mWriteCancelCallCounter.call();
+    }
+
     void onFinishCalled() {
         mFinishCallCounter.call();
     }
@@ -395,6 +402,11 @@
         waitForCallbackCallCount(mWriteCallCounter, count, "Did not get expected call to write.");
     }
 
+    void waitForWriteCancelCallback(int count) {
+        waitForCallbackCallCount(mWriteCancelCallCounter, count,
+                "Did not get expected cancel of write.");
+    }
+
     private static void waitForCallbackCallCount(CallCounter counter, int count, String message) {
         try {
             counter.waitForCount(count, OPERATION_TIMEOUT_MILLIS);
@@ -453,6 +465,7 @@
         mCancelOperationCounter.reset();
         mLayoutCallCounter.reset();
         mWriteCallCounter.reset();
+        mWriteCancelCallCounter.reset();
         mFinishCallCounter.reset();
         mPrintJobQueuedCallCounter.reset();
         mCreateSessionCallCounter.reset();
@@ -619,6 +632,10 @@
         }
     }
 
+    void assertNoPrintButton() throws UiObjectNotFoundException, IOException {
+        assertFalse(sUiDevice.hasObject(By.res("com.android.printspooler:id/print_button")));
+    }
+
     void clickPrintButton() throws UiObjectNotFoundException, IOException {
         try {
             UiObject printButton = sUiDevice.findObject(new UiSelector().resourceId(
diff --git a/tests/tests/print/src/android/print/cts/PrintDocumentAdapterContractTest.java b/tests/tests/print/src/android/print/cts/PrintDocumentAdapterContractTest.java
index 72e0381..fe6fce1 100644
--- a/tests/tests/print/src/android/print/cts/PrintDocumentAdapterContractTest.java
+++ b/tests/tests/print/src/android/print/cts/PrintDocumentAdapterContractTest.java
@@ -236,21 +236,24 @@
     }
 
     @Test
-    public void nonCallingBackWrite() throws Exception {
-        final PrintAttributes[] printAttributes = new PrintAttributes[1];
+    public void nonCallingBackWrite() throws Throwable {
+        final PrintAttributes[] lastLayoutPrintAttributes = new PrintAttributes[1];
         final boolean[] isWriteBroken = new boolean[1];
 
         // Create a mock print adapter.
         PrintDocumentAdapter adapter = createMockPrintDocumentAdapter(
                 invocation -> {
-                    printAttributes[0] = (PrintAttributes) invocation.getArguments()[1];
+                    PrintAttributes printAttributes =
+                            (PrintAttributes) invocation.getArguments()[1];
                     PrintDocumentAdapter.LayoutResultCallback callback =
                             (PrintDocumentAdapter.LayoutResultCallback) invocation
                                     .getArguments()[3];
 
                     callback.onLayoutFinished(new PrintDocumentInfo.Builder(PRINT_JOB_NAME)
                                     .setPageCount(PrintDocumentInfo.PAGE_COUNT_UNKNOWN).build(),
-                            false);
+                            lastLayoutPrintAttributes[0] != printAttributes);
+
+                    lastLayoutPrintAttributes[0] = printAttributes;
 
                     onLayoutCalled();
                     return null;
@@ -264,7 +267,7 @@
                                 .setOnCancelListener(() -> callback.onWriteCancelled());
                     } else {
                         try (ParcelFileDescriptor fd = (ParcelFileDescriptor) args[1]) {
-                            writeBlankPages(printAttributes[0], fd, 0, 1);
+                            writeBlankPages(lastLayoutPrintAttributes[0], fd, 0, 1);
                         }
                         callback.onWriteFinished(new PageRange[]{new PageRange(0, 0)});
                     }
@@ -277,6 +280,79 @@
                 });
 
         // never return from writes until we repair the write call later
+        isWriteBroken[0] = false;
+
+        // Start printing.
+        print(adapter);
+
+        // Wait for write. This will happen as the the first layout always triggers a write
+        waitForWriteAdapterCallback(1);
+
+        // Make write broken
+        isWriteBroken[0] = true;
+        selectPrinter("Second printer");
+
+        assertNoPrintButton();
+
+        // Wait for write (The second printer causes a re-layout as it does not support the default
+        // page size). The write is broken, hence this never returns
+        waitForWriteAdapterCallback(2);
+
+        assertNoPrintButton();
+
+        getUiDevice().pressBack();
+        getUiDevice().pressBack();
+
+        waitForPrinterDiscoverySessionDestroyCallbackCalled(1);
+    }
+
+    @Test
+    public void nonChangePrinterWhileNotWritten() throws Exception {
+        final PrintAttributes[] lastLayoutPrintAttributes = new PrintAttributes[1];
+        final boolean[] isWriteBroken = new boolean[1];
+
+        // Create a mock print adapter.
+        PrintDocumentAdapter adapter = createMockPrintDocumentAdapter(
+                invocation -> {
+                    PrintAttributes printAttributes =
+                            (PrintAttributes) invocation.getArguments()[1];
+                    PrintDocumentAdapter.LayoutResultCallback callback =
+                            (PrintDocumentAdapter.LayoutResultCallback) invocation
+                                    .getArguments()[3];
+
+                    callback.onLayoutFinished(new PrintDocumentInfo.Builder(PRINT_JOB_NAME)
+                                    .setPageCount(PrintDocumentInfo.PAGE_COUNT_UNKNOWN).build(),
+                            lastLayoutPrintAttributes[0] != printAttributes);
+
+                    lastLayoutPrintAttributes[0] = printAttributes;
+
+                    onLayoutCalled();
+                    return null;
+                }, invocation -> {
+                    Object[] args = invocation.getArguments();
+                    PrintDocumentAdapter.WriteResultCallback callback =
+                            (PrintDocumentAdapter.WriteResultCallback) args[3];
+
+                    if (isWriteBroken[0]) {
+                        ((CancellationSignal) args[2]).setOnCancelListener(() -> {
+                            callback.onWriteCancelled();
+                            onWriteCancelCalled();
+                        });
+                    } else {
+                        try (ParcelFileDescriptor fd = (ParcelFileDescriptor) args[1]) {
+                            writeBlankPages(lastLayoutPrintAttributes[0], fd, 0, 1);
+                        }
+                        callback.onWriteFinished(new PageRange[]{new PageRange(0, 1)});
+                    }
+
+                    onWriteCalled();
+                    return null;
+                }, invocation -> {
+                    onFinishCalled();
+                    return null;
+                });
+
+        // Fake a very long write
         isWriteBroken[0] = true;
 
         // Start printing.
@@ -285,17 +361,20 @@
         // Wait for write. This will happen as the the first layout always triggers a write
         waitForWriteAdapterCallback(1);
 
-        // Finally return useful data from the adapter's write call
+        // Repair write and cause a re-layout by changing the printer
         isWriteBroken[0] = false;
+        selectPrinter("Second printer");
 
-        selectPrinter("Third printer");
-        clickPrintButton();
-        answerPrintServicesWarning(true);
-
-        // The layout reports that the doc did not change and none of the print attributes changed.
-        // Still: As the pages are not written yet we still get a write call
+        // The first write should be canceled and the second write succeeds
+        waitForWriteCancelCallback(1);
         waitForWriteAdapterCallback(2);
 
+        // Click the print button.
+        clickPrintButton();
+
+        // Answer the dialog for the print service cloud warning
+        answerPrintServicesWarning(true);
+
         waitForPrinterDiscoverySessionDestroyCallbackCalled(1);
     }