Fix flaky testAddVisualStateCallbackAfterRendererGone

android_webview/javatests/src/org/chromium/android_webview/test/crash/VisualStateCallbackTest.java

Index: android_webview/javatests/src/org/chromium/android_webview/test/crash/VisualStateCallbackTest.java
diff --git a/android_webview/javatests/src/org/chromium/android_webview/test/crash/VisualStateCallbackTest.java b/android_webview/javatests/src/org/chromium/android_webview/test/crash/VisualStateCallbackTest.java
index 447dd49c303e84b90f7bec0e382fd365bc7b576e..884f4c0c81305c0d8786a3fce892f975984824a5 100644
--- a/android_webview/javatests/src/org/chromium/android_webview/test/crash/VisualStateCallbackTest.java
+++ b/android_webview/javatests/src/org/chromium/android_webview/test/crash/VisualStateCallbackTest.java
@@ -153,13 +153,16 @@ public class VisualStateCallbackTest extends AwTestBase {
     @SmallTest
     @CommandLineFlags
             .Add(AwSwitches.WEBVIEW_SANDBOXED_RENDERER)
-            @ParameterizedTest.Set
-            public void testAddVisualStateCallbackAfterRendererGone() throws Throwable {
+    @ParameterizedTest.Set
+    public void testAddVisualStateCallbackAfterRendererGone() throws Throwable {
+        final VisualStateCallbackImpl vsImpl = new VisualStateCallbackImpl();
+        mHelper.setOnRenderProcessGoneTask(new Runnable() {
+            @Override
+            public void run() {
+                mAwContents.insertVisualStateCallback(vsImpl.requestId(), vsImpl);
+            }
+        });
         loadUrlAsync(mAwContents, "chrome://kill");
-        mHelper.waitForRenderProcessGone();
-
-        VisualStateCallbackImpl vsImpl = new VisualStateCallbackImpl();
-        insertVisualStateCallbackOnUIThread(mAwContents, vsImpl.requestId(), vsImpl);
 
         mHelper.waitForRenderProcessGoneNotifiedToAwContentsClient();
 
@@ -174,8 +177,8 @@ public class VisualStateCallbackTest extends AwTestBase {
     @SmallTest
     @CommandLineFlags
             .Add(AwSwitches.WEBVIEW_SANDBOXED_RENDERER)
-            @ParameterizedTest.Set
-            public void testVisualStateCallbackNotCalledAfterRendererGone() throws Throwable {
+    @ParameterizedTest.Set
+    public void testVisualStateCallbackNotCalledAfterRendererGone() throws Throwable {
 
         VisualStateCallbackImpl vsImpl = new VisualStateCallbackImpl();
         insertVisualStateCallbackOnUIThread(mAwContents, vsImpl.requestId(), vsImpl);