[WebView] Add unique visual state request ids.

android_webview/javatests/src/org/chromium/android_webview/test/VisualStateTest.java

Index: android_webview/javatests/src/org/chromium/android_webview/test/VisualStateTest.java
diff --git a/android_webview/javatests/src/org/chromium/android_webview/test/VisualStateTest.java b/android_webview/javatests/src/org/chromium/android_webview/test/VisualStateTest.java
index a20ddb5f79f5cd14367042c38f9a22a2c0bfb6bc..7b06c16a10cea29e9a710bed631e6ab0b0476acb 100644
--- a/android_webview/javatests/src/org/chromium/android_webview/test/VisualStateTest.java
+++ b/android_webview/javatests/src/org/chromium/android_webview/test/VisualStateTest.java
@@ -18,7 +18,6 @@ import org.chromium.content_public.browser.LoadUrlParams;
 
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.atomic.AtomicReference;
 
 /**
  * Visual state related tests.
@@ -26,6 +25,8 @@ import java.util.concurrent.atomic.AtomicReference;
 public class VisualStateTest extends AwTestBase {
 
     private TestAwContentsClient mContentsClient = new TestAwContentsClient();
+    private long mExpectedFlushRequestId = -1;
+    private AwContents mAwContents;
 
     @Feature({"AndroidWebView"})
     @SmallTest
@@ -39,17 +40,19 @@ public class VisualStateTest extends AwTestBase {
         runTestOnUiThread(new Runnable() {
             @Override
             public void run() {
-                awContents.flushVisualState(new AwContents.VisualStateFlushCallback() {
-                    @Override
-                    public void onComplete() {
-                        ch.notifyCalled();
-                    }
+                mExpectedFlushRequestId =
+                        awContents.flushVisualState(new AwContents.VisualStateFlushCallback() {
+                            @Override
+                            public void onComplete(long requestId) {
+                                assertEquals(mExpectedFlushRequestId, requestId);
+                                ch.notifyCalled();
+                            }
 
-                    @Override
-                    public void onFailure() {
-                        fail("onFailure received");
-                    }
-                });
+                            @Override
+                            public void onFailure(long requestId) {
+                                fail("onFailure received");
+                            }
+                        });
             }
         });
         ch.waitForCallback(chCount);
@@ -65,44 +68,44 @@ public class VisualStateTest extends AwTestBase {
         final LoadUrlParams bluePageUrl = createTestPageUrl("blue");
         final CountDownLatch testFinishedSignal = new CountDownLatch(1);
 
-        final AtomicReference<AwContents> awContentsRef = new AtomicReference<>();
         final AwTestContainerView testView =
                 createAwTestContainerViewOnMainSync(new TestAwContentsClient() {
                     @Override
                     public void onPageFinished(String url) {
                         if (bluePageUrl.getUrl().equals(url)) {
-                            awContentsRef.get().flushVisualState(new VisualStateFlushCallback() {
-                                @Override
-                                public void onFailure() {
-                                    fail("onFailure received");
-                                }
+                            mExpectedFlushRequestId =
+                                    mAwContents.flushVisualState(new VisualStateFlushCallback() {
+                                        @Override
+                                        public void onFailure(long requestId) {
+                                            fail("onFailure received");
+                                        }
 
-                                @Override
-                                public void onComplete() {
-                                    Bitmap blueScreenshot = GraphicsTestUtils.drawAwContents(
-                                            awContentsRef.get(), 1, 1);
-                                    assertEquals(Color.BLUE, blueScreenshot.getPixel(0, 0));
-                                    testFinishedSignal.countDown();
-                                }
-                            });
+                                        @Override
+                                        public void onComplete(long requestId) {
+                                            assertEquals(mExpectedFlushRequestId, requestId);
+                                            Bitmap blueScreenshot =
+                                                    GraphicsTestUtils.drawAwContents(
+                                                            mAwContents, 1, 1);
+                                            assertEquals(Color.BLUE, blueScreenshot.getPixel(0, 0));
+                                            testFinishedSignal.countDown();
+                                        }
+                                    });
                         }
                     }
                 });
-        final AwContents awContents = testView.getAwContents();
-        awContentsRef.set(awContents);
+        mAwContents = testView.getAwContents();
 
         runTestOnUiThread(new Runnable() {
             @Override
             public void run() {
-                awContents.setBackgroundColor(Color.RED);
-                awContents.loadUrl(bluePageUrl);
+                mAwContents.setBackgroundColor(Color.RED);
+                mAwContents.loadUrl(bluePageUrl);
 
                 // We have just loaded the blue page, but the graphics pipeline is asynchronous
                 // so at this point the WebView still draws red, ie. the View background color.
                 // Only when the flush callback is received will we know for certain that the
                 // blue page contents are on screen.
-                Bitmap redScreenshot = GraphicsTestUtils.drawAwContents(
-                        awContentsRef.get(), 1, 1);
+                Bitmap redScreenshot = GraphicsTestUtils.drawAwContents(mAwContents, 1, 1);
                 assertEquals(Color.RED, redScreenshot.getPixel(0, 0));
             }
         });