[Paint Preview] Improve PaintPreviewTabService tests

Improve test coverage for PaintPreviewTabService by covering
auditing.

Change-Id: I4482732105561f1576900a79f4d68479b90f4604
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2402161Reviewed-by: Mehran Mahmoudi <mahmoudi@chromium.org>
Commit-Queue: Calder Kitagawa <ckitagawa@chromium.org>
Cr-Commit-Position: refs/heads/master@{#805770}

chrome/android/javatests/src/org/chromium/chrome/browser/paint_preview/services/PaintPreviewTabServiceTest.java

@@ -150,6 +150,44 @@ public class PaintPreviewTabServiceTest {
         }, "Paint Preview didn't get deleted.", TIMEOUT_MS, POLLING_INTERVAL_MS);
     }
 
+    /**
+     * Tests that capturing and deleting via an audit works as expected.
+     */
+    @Test
+    @MediumTest
+    @Feature({"PaintPreview"})
+    public void testCapturedAndDeleteViaAudit() throws Exception {
+        EmbeddedTestServer testServer = mActivityTestRule.getTestServer();
+        final String url = testServer.getURL("/chrome/test/data/android/about.html");
+
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            mPaintPreviewTabService = PaintPreviewTabServiceFactory.getServiceInstance();
+            // Use capture on switch mode.
+            mPaintPreviewTabService.onRestoreCompleted(mTabModelSelector, true, true);
+            mTab.loadUrl(new LoadUrlParams(url));
+        });
+        // Give the tab time to complete layout before hiding.
+        TimeUnit.SECONDS.sleep(1);
+
+        // This will hide mTab so that a capture occurs.
+        mActivityTestRule.loadUrlInNewTab(url);
+
+        int tabId = mTab.getId();
+        CriteriaHelper.pollUiThread(() -> {
+            return mPaintPreviewTabService.hasCaptureForTab(tabId);
+        }, "Paint Preview didn't get captured.", TIMEOUT_MS, POLLING_INTERVAL_MS);
+
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            // Use the incognito tab model as the normal tab model will still have the tab ids
+            // active.
+            mPaintPreviewTabService.auditOnStart(mTabModelSelector.getModel(/*incognito=*/true));
+        });
+
+        CriteriaHelper.pollUiThread(() -> {
+            return !mPaintPreviewTabService.hasCaptureForTab(tabId);
+        }, "Paint Preview didn't get deleted.", TIMEOUT_MS, POLLING_INTERVAL_MS);
+    }
+
     /**
      * Verifies the early cache is created correctly.
      */

chrome/browser/paint_preview/android/java/src/org/chromium/chrome/browser/paint_preview/services/PaintPreviewTabService.java

@@ -142,7 +142,8 @@ public class PaintPreviewTabService implements NativePaintPreviewServiceProvider
                 AUDIT_START_DELAY_MS);
     }
 
-    private void auditOnStart(TabModel regularTabModel) {
+    @VisibleForTesting
+    void auditOnStart(TabModel regularTabModel) {
         int tabCount = regularTabModel.getCount();
         int[] tabIds = new int[tabCount];
         for (int i = 0; i < tabCount; i++) {