Introduce ThreadedInputConnection behind a switch

content/public/android/javatests/src/org/chromium/content/browser/input/ImeTest.java

Index: content/public/android/javatests/src/org/chromium/content/browser/input/ImeTest.java
diff --git a/content/public/android/javatests/src/org/chromium/content/browser/input/ImeTest.java b/content/public/android/javatests/src/org/chromium/content/browser/input/ImeTest.java
index 704441b743a5ac6465a3e649ad94fb057d1d0aa8..7cd2eeb5793bca16ece4316bf0fee913822b2a81 100644
--- a/content/public/android/javatests/src/org/chromium/content/browser/input/ImeTest.java
+++ b/content/public/android/javatests/src/org/chromium/content/browser/input/ImeTest.java
@@ -9,13 +9,18 @@ import android.content.ClipData;
 import android.content.ClipboardManager;
 import android.content.Context;
 import android.content.res.Configuration;
+import android.os.Handler;
+import android.os.Looper;
 import android.test.suitebuilder.annotation.MediumTest;
 import android.test.suitebuilder.annotation.SmallTest;
 import android.text.TextUtils;
+import android.util.Pair;
 import android.view.KeyEvent;
 import android.view.View;
 import android.view.inputmethod.EditorInfo;
+import android.view.inputmethod.InputConnection;
 
+import org.chromium.base.ObserverList;
 import org.chromium.base.ThreadUtils;
 import org.chromium.base.test.util.Feature;
 import org.chromium.base.test.util.UrlUtils;
@@ -25,8 +30,8 @@ import org.chromium.content.browser.test.util.CriteriaHelper;
 import org.chromium.content.browser.test.util.DOMUtils;
 import org.chromium.content.browser.test.util.TestCallbackHelperContainer;
 import org.chromium.content.browser.test.util.TestInputMethodManagerWrapper;
-import org.chromium.content.browser.test.util.TestInputMethodManagerWrapper.Range;
 import org.chromium.content_public.browser.WebContents;
+import org.chromium.content_public.browser.WebContentsObserver;
 import org.chromium.content_shell_apk.ContentShellTestBase;
 import org.chromium.ui.base.ime.TextInputType;
 
@@ -55,8 +60,8 @@ public class ImeTest extends ContentShellTestBase {
             + "<br/><p><span id=\"plain_text\">This is Plain Text One</span></p>"
             + "</form></body></html>");
 
-    private TestAdapterInputConnection mConnection;
-    private TestAdapterInputConnectionFactory mConnectionFactory;
+    private ChromiumBaseInputConnection mConnection;
+    private TestInputConnectionFactory mConnectionFactory;
     private ImeAdapter mImeAdapter;
 
     private ContentViewCore mContentViewCore;
@@ -76,7 +81,8 @@ public class ImeTest extends ContentShellTestBase {
         mInputMethodManagerWrapper = new TestInputMethodManagerWrapper(mContentViewCore);
         getImeAdapter().setInputMethodManagerWrapperForTest(mInputMethodManagerWrapper);
         assertEquals(0, mInputMethodManagerWrapper.getShowSoftInputCounter());
-        mConnectionFactory = new TestAdapterInputConnectionFactory();
+        mConnectionFactory =
+                new TestInputConnectionFactory(getImeAdapter().getInputConnectionFactoryForTest());
         getImeAdapter().setInputConnectionFactory(mConnectionFactory);
 
         mCallbackContainer = new TestCallbackHelperContainer(mContentViewCore);
@@ -86,46 +92,68 @@ public class ImeTest extends ContentShellTestBase {
         DOMUtils.clickNode(this, mContentViewCore, "input_text");
         assertWaitForKeyboardStatus(true);
 
-        mConnection = (TestAdapterInputConnection) getAdapterInputConnection();
+        mConnection = getInputConnection();
         mImeAdapter = getImeAdapter();
 
-        waitAndVerifyStatesAndCalls(0, "", 0, 0, -1, -1);
+        if (usingReplicaInputConnection()) {
+            // This is not needed if onCreateInputConnection() can return correct selection range.
+            waitAndVerifyUpdateSelection(0, 0, 0, -1, -1);
+        }
         waitForKeyboardStates(1, 0, 1, new Integer[] {TextInputType.TEXT});
-        assertEquals(0, mInputMethodManagerWrapper.getEditorInfo().initialSelStart);
-        assertEquals(0, mInputMethodManagerWrapper.getEditorInfo().initialSelEnd);
+        assertEquals(0, mConnectionFactory.getOutAttrs().initialSelStart);
+        assertEquals(0, mConnectionFactory.getOutAttrs().initialSelEnd);
 
         resetAllStates();
     }
 
-    private void assertNoFurtherStateUpdate(final int index) throws InterruptedException {
-        final List<TestImeState> states = mConnectionFactory.getImeStateList();
-        try {
-            CriteriaHelper.pollForUIThreadCriteria(new Criteria() {
-                @Override
-                public boolean isSatisfied() {
-                    return states.size() > index;
-                }
-            });
-            fail("Unexpected updates pending");
-        } catch (AssertionError e) {
-            // TODO(tedchoc): This is horrible and should never timeout to determine success.
-        }
+    @MediumTest
+    @Feature({"TextInput", "Main"})
+    public void testKeyboardDismissedAfterClickingGo() throws Throwable {
+        setComposingText("hello", 1);
+        waitAndVerifyUpdateSelection(0, 5, 5, 0, 5);
+
+        performGo(mCallbackContainer);
+
+        assertWaitForKeyboardStatus(false);
     }
 
     @MediumTest
     @Feature({"TextInput", "Main"})
-    public void testSetUpGeneratesNoFurtherStateUpdate() throws Throwable {
-        assertNoFurtherStateUpdate(0);
-        waitForKeyboardStates(0, 0, 0, new Integer[] {});
+    public void testDoesNotHang_getTextAfterKeyboardHides() throws Throwable {
+        setComposingText("hello", 1);
+        waitAndVerifyUpdateSelection(0, 5, 5, 0, 5);
+
+        performGo(mCallbackContainer);
+
+        // This may time out if we do not get the information on time.
+        for (int i = 0; i < 100; ++i) {
+            getTextBeforeCursor(10, 0);
+        }
+
+        assertWaitForKeyboardStatus(false);
     }
 
     @MediumTest
     @Feature({"TextInput", "Main"})
-    public void testKeyboardDismissedAfterClickingGo() throws Throwable {
+    public void testDoesNotHang_rendererCrashes() throws Throwable {

[no sievers, 2016/02/25 01:41:44]

per chat: can you remove this one and the test above, since they look like they
don't actually test anything (and the for-loop is odd)?
and looks like it wouldn't cause a hard deadlock for the user anyways. one can
always still interact with the UI, and showing the sad page actually takes care
of removing the focus and unblocking in a way.

[Changwan Ryu, 2016/02/25 01:56:12]

Removed this test as it does not test anything.
I still think that the other test has some value (it tests something), so I'd
like to keep it.
I've left a TODO to remove the odd loop there.

         setComposingText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, 0, 5);
+        waitAndVerifyUpdateSelection(0, 5, 5, 0, 5);
 
-        performGo(mCallbackContainer);
+        ThreadUtils.runOnUiThread(new Runnable() {
+            @Override
+            public void run() {
+                ObserverList.RewindableIterator<WebContentsObserver> observers =
+                        getWebContents().getObserversForTesting();
+                while (observers.hasNext()) {
+                    observers.next().renderProcessGone(false);
+                }
+            }
+        });
+
+        // This may time out if we do not get the information on time.
+        for (int i = 0; i < 100; ++i) {
+            getTextBeforeCursor(10, 0);
+        }
 
         assertWaitForKeyboardStatus(false);
     }
@@ -134,16 +162,24 @@ public class ImeTest extends ContentShellTestBase {
     @Feature({"TextInput", "Main"})
     public void testCommitWhileComposingText() throws Throwable {
         setComposingText("h", 1);
-        waitAndVerifyStates(0, "h", 1, 1, 0, 1);
+        waitAndVerifyUpdateSelection(0, 1, 1, 0, 1);
 
         setComposingText("he", 1);
-        waitAndVerifyStates(1, "he", 2, 2, 0, 2);
+        waitAndVerifyUpdateSelection(1, 2, 2, 0, 2);
 
         setComposingText("hel", 1);
-        waitAndVerifyStates(2, "hel", 3, 3, 0, 3);
+        waitAndVerifyUpdateSelection(2, 3, 3, 0, 3);
 
         commitText("hel", 1);
-        waitAndVerifyStates(3, "hel", 3, 3, -1, -1);
+        waitAndVerifyUpdateSelection(3, 3, 3, -1, -1);
+
+        setComposingText("lo", 1);
+        waitAndVerifyUpdateSelection(4, 5, 5, 3, 5);
+
+        commitText("", 1);
+        waitAndVerifyUpdateSelection(5, 3, 3, -1, -1);
+
+        assertTextsAroundCursor("hel", "", "");
     }
 
     @SmallTest
@@ -152,26 +188,28 @@ public class ImeTest extends ContentShellTestBase {
         focusElementAndWaitForStateUpdate("textarea");
 
         setComposingText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, 0, 5);
+        waitAndVerifyUpdateSelection(0, 5, 5, 0, 5);
 
         // Cancel the current composition and replace it with enter.
         commitText("\n", 1);
+        waitAndVerifyUpdateSelection(1, 1, 1, -1, -1);
         // The second new line is not a user visible/editable one, it is a side-effect of Blink
         // using <br> internally. This only happens when \n is at the end.
-        waitAndVerifyStatesAndCalls(1, "\n\n", 1, 1, -1, -1);
+        assertTextsAroundCursor("\n", "", "\n");
 
         commitText("world", 1);
-        waitAndVerifyStatesAndCalls(2, "\nworld", 6, 6, -1, -1);
+        waitAndVerifyUpdateSelection(2, 6, 6, -1, -1);
+        assertTextsAroundCursor("\nworld", "", "");
     }
 
     @SmallTest
     @Feature({"TextInput"})
     public void testImeCopy() throws Exception {
         commitText("hello", 1);
-        waitAndVerifyStates(0, "hello", 5, 5, -1, -1);
+        waitAndVerifyUpdateSelection(0, 5, 5, -1, -1);
 
         setSelection(2, 5);
-        waitAndVerifyStates(1, "hello", 2, 5, -1, -1);
+        waitAndVerifyUpdateSelection(1, 2, 5, -1, -1);
 
         copy();
         assertClipboardContents(getActivity(), "llo");
@@ -181,16 +219,13 @@ public class ImeTest extends ContentShellTestBase {
     @Feature({"TextInput"})
     public void testEnterTextAndRefocus() throws Exception {
         commitText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, -1, -1);
-
-        DOMUtils.clickNode(this, mContentViewCore, "input_radio");
-        assertWaitForKeyboardStatus(false);
-
+        waitAndVerifyUpdateSelection(0, 5, 5, -1, -1);
+        restartInput();
         DOMUtils.clickNode(this, mContentViewCore, "input_text");
         assertWaitForKeyboardStatus(true);
 
-        assertEquals(5, mInputMethodManagerWrapper.getEditorInfo().initialSelStart);
-        assertEquals(5, mInputMethodManagerWrapper.getEditorInfo().initialSelEnd);
+        assertEquals(5, mConnectionFactory.getOutAttrs().initialSelStart);
+        assertEquals(5, mConnectionFactory.getOutAttrs().initialSelEnd);
     }
 
     @SmallTest
@@ -203,9 +238,16 @@ public class ImeTest extends ContentShellTestBase {
 
         // When input connection is null, we still need to set flags to prevent InputMethodService
         // from entering fullscreen mode and from opening custom UI.
-        assertNull(mInputMethodManagerWrapper.getInputConnection());
-        assertEquals(EditorInfo.IME_FLAG_NO_FULLSCREEN | EditorInfo.IME_FLAG_NO_EXTRACT_UI,
-                mInputMethodManagerWrapper.getEditorInfo().imeOptions);
+        CriteriaHelper.pollForUIThreadCriteria(new Criteria() {
+            @Override
+            public boolean isSatisfied() {
+                return getInputConnection() == null;
+            }
+        });
+        assertTrue(
+                (mConnectionFactory.getOutAttrs().imeOptions
+                        & (EditorInfo.IME_FLAG_NO_FULLSCREEN | EditorInfo.IME_FLAG_NO_EXTRACT_UI))
+                != 0);
 
         // showSoftInput(), restartInput()
         focusElement("input_number1");
@@ -232,6 +274,27 @@ public class ImeTest extends ContentShellTestBase {
                 TextInputType.TEXT});
     }
 
+    private void assertTextsAroundCursor(
+            CharSequence before, CharSequence selected, CharSequence after) throws Exception {
+        assertEquals(before, getTextBeforeCursor(100, 0));
+
+        CharSequence actualSelected = getSelectedText(0);
+        if (usingReplicaInputConnection() && TextUtils.isEmpty(actualSelected)) {
+            // ReplicaInputConnection will return null but ChromiumInputConnection will return "".
+            actualSelected = "";
+        }
+        assertEquals(selected, actualSelected);
+
+        if (usingReplicaInputConnection() && after.equals("\n")) {
+            // When the text ends with \n, we have a second new line that is not user
+            // visible/editable one, it is a side effect of using <br> internally.
+            // Replica model simply deviates from the blink editor in this case.
+            assertEquals("", getTextAfterCursor(100, 0));
+        } else {
+            assertEquals(after, getTextAfterCursor(100, 0));
+        }
+    }
+
     private void waitForKeyboardStates(int show, int hide, int restart, Integer[] history)
             throws InterruptedException {
         final String expected = stringifyKeyboardStates(show, hide, restart, history);
@@ -246,27 +309,16 @@ public class ImeTest extends ContentShellTestBase {
     }
 
     private void resetAllStates() {
-        mInputMethodManagerWrapper.resetCounters();
+        mInputMethodManagerWrapper.reset();
         mConnectionFactory.clearTextInputTypeHistory();
-        resetUpdateStateList();
     }
 
     private String getKeyboardStates() {
-        try {
-            return ThreadUtils.runOnUiThreadBlocking(new Callable<String>() {
-                @Override
-                public String call() throws Exception {
-                    int showCount = mInputMethodManagerWrapper.getShowSoftInputCounter();
-                    int hideCount = mInputMethodManagerWrapper.getHideSoftInputCounter();
-                    int restartCount = mInputMethodManagerWrapper.getRestartInputCounter();
-                    Integer[] history = mConnectionFactory.getTextInputTypeHistory();
-                    return stringifyKeyboardStates(showCount, hideCount, restartCount, history);
-                }
-            });
-        } catch (ExecutionException e) {
-            e.printStackTrace();
-            return null;
-        }
+        int showCount = mInputMethodManagerWrapper.getShowSoftInputCounter();
+        int hideCount = mInputMethodManagerWrapper.getHideSoftInputCounter();
+        int restartCount = mInputMethodManagerWrapper.getRestartInputCounter();
+        Integer[] history = mConnectionFactory.getTextInputTypeHistory();
+        return stringifyKeyboardStates(showCount, hideCount, restartCount, history);
     }
 
     private String stringifyKeyboardStates(int show, int hide, int restart, Integer[] history) {
@@ -278,7 +330,7 @@ public class ImeTest extends ContentShellTestBase {
     @Feature({"TextInput"})
     public void testKeyboardNotDismissedAfterCopySelection() throws Exception {
         commitText("Sample Text", 1);
-        waitAndVerifyStatesAndCalls(0, "Sample Text", 11, 11, -1, -1);
+        waitAndVerifyUpdateSelection(0, 11, 11, -1, -1);
 
         // This will select 'Text' part.
         DOMUtils.clickNode(this, mContentViewCore, "input_text");
@@ -299,7 +351,7 @@ public class ImeTest extends ContentShellTestBase {
     @Feature({"TextInput"})
     public void testImeNotDismissedAfterCutSelection() throws Exception {
         commitText("Sample Text", 1);
-        waitAndVerifyStatesAndCalls(0, "Sample Text", 11, 11, -1, -1);
+        waitAndVerifyUpdateSelection(0, 11, 11, -1, -1);
         DOMUtils.longPressNode(this, mContentViewCore, "input_text");
         assertWaitForSelectActionBarStatus(true);
         assertWaitForKeyboardStatus(true);
@@ -334,20 +386,20 @@ public class ImeTest extends ContentShellTestBase {
     public void testLongPressInputWhileComposingText() throws Exception {
         assertWaitForSelectActionBarStatus(false);
         setComposingText("Sample Text", 1);
-        waitAndVerifyStatesAndCalls(0, "Sample Text", 11, 11, 0, 11);
+        waitAndVerifyUpdateSelection(0, 11, 11, 0, 11);
         DOMUtils.longPressNode(this, mContentViewCore, "input_text");
 
         assertWaitForSelectActionBarStatus(true);
 
         // Long press will first change selection region, and then trigger IME app to show up.
         // See RenderFrameImpl::didChangeSelection() and RenderWidget::didHandleGestureEvent().
-        waitAndVerifyStatesAndCalls(1, "Sample Text", 7, 11, 0, 11);
+        waitAndVerifyUpdateSelection(1, 7, 11, 0, 11);
 
         // Now IME app wants to finish composing text because an external selection
         // change has been detected. At least Google Latin IME and Samsung IME
         // behave this way.
         finishComposingText();
-        waitAndVerifyStatesAndCalls(2, "Sample Text", 7, 11, -1, -1);
+        waitAndVerifyUpdateSelection(2, 7, 11, -1, -1);
     }
 
     @SmallTest
@@ -493,14 +545,14 @@ public class ImeTest extends ContentShellTestBase {
     @Feature({"TextInput"})
     public void testImeCut() throws Exception {
         commitText("snarful", 1);
-        waitAndVerifyStatesAndCalls(0, "snarful", 7, 7, -1, -1);
+        waitAndVerifyUpdateSelection(0, 7, 7, -1, -1);
 
         setSelection(1, 5);
-        waitAndVerifyStatesAndCalls(1, "snarful", 1, 5, -1, -1);
+        waitAndVerifyUpdateSelection(1, 1, 5, -1, -1);
 
         cut();
-        waitAndVerifyStatesAndCalls(2, "sul", 1, 1, -1, -1);
-
+        waitAndVerifyUpdateSelection(2, 1, 1, -1, -1);
+        assertTextsAroundCursor("s", "", "ul");
         assertClipboardContents(getActivity(), "narf");
     }
 
@@ -518,28 +570,33 @@ public class ImeTest extends ContentShellTestBase {
         });
 
         paste();
-        waitAndVerifyStatesAndCalls(0, "blarg", 5, 5, -1, -1);
+        // Paste is a two step process when there is a non-zero selection.
+        waitAndVerifyUpdateSelection(0, 5, 5, -1, -1);
+        assertTextsAroundCursor("blarg", "", "");
 
         setSelection(3, 5);
-        waitAndVerifyStatesAndCalls(1, "blarg", 3, 5, -1, -1);
+        waitAndVerifyUpdateSelection(1, 3, 5, -1, -1);
+        assertTextsAroundCursor("bla", "rg", "");
 
         paste();
         // Paste is a two step process when there is a non-zero selection.
-        waitAndVerifyStates(2, "bla", 3, 3, -1, -1);
-        waitAndVerifyStatesAndCalls(3, "blablarg", 8, 8, -1, -1);
+        waitAndVerifyUpdateSelection(2, 3, 3, -1, -1);
+        waitAndVerifyUpdateSelection(3, 8, 8, -1, -1);
+        assertTextsAroundCursor("blablarg", "", "");
 
         paste();
-        waitAndVerifyStatesAndCalls(4, "blablargblarg", 13, 13, -1, -1);
+        waitAndVerifyUpdateSelection(4, 13, 13, -1, -1);
+        assertTextsAroundCursor("blablargblarg", "", "");
     }
 
     @SmallTest
     @Feature({"TextInput"})
     public void testImeSelectAndUnSelectAll() throws Exception {
         commitText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, -1, -1);
+        waitAndVerifyUpdateSelection(0, 5, 5, -1, -1);
 
         selectAll();
-        waitAndVerifyStatesAndCalls(1, "hello", 0, 5, -1, -1);
+        waitAndVerifyUpdateSelection(1, 0, 5, -1, -1);
 
         unselect();
 
@@ -564,83 +621,75 @@ public class ImeTest extends ContentShellTestBase {
         focusElementAndWaitForStateUpdate("textarea");
 
         commitText("hllo", 1);
-        waitAndVerifyStatesAndCalls(0, "hllo", 4, 4, -1, -1);
+        waitAndVerifyUpdateSelection(0, 4, 4, -1, -1);
 
         commitText(" ", 1);
-        waitAndVerifyStatesAndCalls(1, "hllo ", 5, 5, -1, -1);
+        waitAndVerifyUpdateSelection(1, 5, 5, -1, -1);
 
         setSelection(1, 1);
-        waitAndVerifyStatesAndCalls(2, "hllo ", 1, 1, -1, -1);
+        waitAndVerifyUpdateSelection(2, 1, 1, -1, -1);
+        assertTextsAroundCursor("h", "", "llo ");
 
         setComposingRegion(0, 4);
-        waitAndVerifyStatesAndCalls(3, "hllo ", 1, 1, 0, 4);
+        waitAndVerifyUpdateSelection(3, 1, 1, 0, 4);
 
         finishComposingText();
-        waitAndVerifyStatesAndCalls(4, "hllo ", 1, 1, -1, -1);
+        waitAndVerifyUpdateSelection(4, 1, 1, -1, -1);
 
         commitText("\n", 1);
-        waitAndVerifyStatesAndCalls(5, "h\nllo ", 2, 2, -1, -1);
+        waitAndVerifyUpdateSelection(5, 2, 2, -1, -1);
+        assertTextsAroundCursor("h\n", "", "llo ");
     }
 
-    /*
+    // http://crbug.com/445499
     @SmallTest
     @Feature({"TextInput", "Main"})
-    http://crbug.com/445499
-    */
     public void testDeleteText() throws Throwable {
         focusElement("textarea");
 
         // The calls below are a reflection of what the stock Google Keyboard (Andr
         // when the noted key is touched on screen.
         // H
-        resetUpdateStateList();
         setComposingText("h", 1);
-        assertUpdateStateCall(1000);
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("h", getTextBeforeCursor(9, 0));
 
         // O
-        resetUpdateStateList();
         setComposingText("ho", 1);
-        assertUpdateStateCall(1000);
-        assertEquals("ho", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("ho", getTextBeforeCursor(9, 0));
 
-        resetUpdateStateList();
         setComposingText("h", 1);
-        assertUpdateStateCall(1000);
         setComposingRegion(0, 1);
         setComposingText("h", 1);
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("h", getTextBeforeCursor(9, 0));
 
         // I
         setComposingText("hi", 1);
-        assertEquals("hi", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("hi", getTextBeforeCursor(9, 0));
 
         // SPACE
         commitText("hi", 1);
         commitText(" ", 1);
-        assertEquals("hi ", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("hi ", getTextBeforeCursor(9, 0));
 
         // DEL
         deleteSurroundingText(1, 0);
         setComposingRegion(0, 2);
-        assertEquals("hi", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("hi", getTextBeforeCursor(9, 0));
 
         setComposingText("h", 1);
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("h", getTextBeforeCursor(9, 0));
 
         commitText("", 1);
-        assertEquals("", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("", getTextBeforeCursor(9, 0));
 
         // DEL (on empty input)
         deleteSurroundingText(1, 0); // DEL on empty still sends 1,0
-        assertEquals("", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("", getTextBeforeCursor(9, 0));
     }
 
 
-    /*
     @SmallTest
     @Feature({"TextInput", "Main"})
-    */
     public void testSwipingText() throws Throwable {
         focusElement("textarea");
 
@@ -649,26 +698,20 @@ public class ImeTest extends ContentShellTestBase {
         // that the test reflects reality.  If this test breaks, it's possible that code has
         // changed and different calls need to be made instead.
         // "three"
-        resetUpdateStateList();
         setComposingText("three", 1);
-        assertUpdateStateCall(1000);
-        assertEquals("three", mConnection.getTextBeforeCursor(99, 0));
+        assertEquals("three", getTextBeforeCursor(99, 0));
 
         // "word"
         commitText("three", 1);
         commitText(" ", 1);
         setComposingText("word", 1);
-        resetUpdateStateList();
-        assertUpdateStateCall(1000);
-        assertEquals("three word", mConnection.getTextBeforeCursor(99, 0));
+        assertEquals("three word", getTextBeforeCursor(99, 0));
 
         // "test"
         commitText("word", 1);
         commitText(" ", 1);
-        resetUpdateStateList();
         setComposingText("test", 1);
-        assertUpdateStateCall(1000);
-        assertEquals("three word test", mConnection.getTextBeforeCursor(99, 0));
+        assertEquals("three word test", getTextBeforeCursor(99, 0));
     }
 
     @SmallTest
@@ -678,19 +721,20 @@ public class ImeTest extends ContentShellTestBase {
         final String smiley = "\uD83D\uDE0A";
 
         commitText(smiley, 1);
-        waitAndVerifyStatesAndCalls(0, smiley, 2, 2, -1, -1);
+        waitAndVerifyUpdateSelection(0, 2, 2, -1, -1);
+        assertTextsAroundCursor(smiley, "", "");
 
         // DEL, sent via dispatchKeyEvent like it is in Android WebView or a physical keyboard.
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_DEL));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_DEL));
 
-        waitAndVerifyStatesAndCalls(1, "", 0, 0, -1, -1);
+        waitAndVerifyUpdateSelection(1, 0, 0, -1, -1);
 
         // Make sure that we accept further typing after deleting the smiley.
         setComposingText("s", 1);
-        waitAndVerifyStatesAndCalls(2, "s", 1, 1, 0, 1);
         setComposingText("sm", 1);
-        waitAndVerifyStatesAndCalls(3, "sm", 2, 2, 0, 2);
+        waitAndVerifyUpdateSelection(2, 1, 1, 0, 1);
+        waitAndVerifyUpdateSelection(3, 2, 2, 0, 2);
     }
 
     @SmallTest
@@ -699,28 +743,19 @@ public class ImeTest extends ContentShellTestBase {
         focusElement("textarea");
 
         // H
-        resetUpdateStateList();
         commitText("h", 1);
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
-        assertUpdateStateCall(1000);
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("h", getTextBeforeCursor(9, 0));
 
         // O
-        resetUpdateStateList();
         commitText("o", 1);
-        assertEquals("ho", mConnection.getTextBeforeCursor(9, 0));
-        assertUpdateStateCall(1000);
-        assertEquals("ho", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("ho", getTextBeforeCursor(9, 0));
 
         // DEL, sent via dispatchKeyEvent like it is in Android WebView or a physical keyboard.
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_DEL));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_DEL));
 
         // DEL
-        resetUpdateStateList();
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
-        assertUpdateStateCall(1000);
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("h", getTextBeforeCursor(9, 0));
     }
 
     @SmallTest
@@ -729,18 +764,12 @@ public class ImeTest extends ContentShellTestBase {
         focusElement("textarea");
 
         // H
-        resetUpdateStateList();
         commitText("h", 1);
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
-        assertUpdateStateCall(1000);
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("h", getTextBeforeCursor(9, 0));
 
         // O
-        resetUpdateStateList();
         commitText("o", 1);
-        assertEquals("ho", mConnection.getTextBeforeCursor(9, 0));
-        assertUpdateStateCall(1000);
-        assertEquals("ho", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("ho", getTextBeforeCursor(9, 0));
 
         // Multiple keydowns should each delete one character (this is for physical keyboard
         // key-repeat).
@@ -749,10 +778,7 @@ public class ImeTest extends ContentShellTestBase {
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_DEL));
 
         // DEL
-        resetUpdateStateList();
-        assertEquals("", mConnection.getTextBeforeCursor(9, 0));
-        assertUpdateStateCall(1000);
-        assertEquals("", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("", getTextBeforeCursor(9, 0));
     }
 
     @SmallTest
@@ -763,17 +789,19 @@ public class ImeTest extends ContentShellTestBase {
         // Type 'a' using a physical keyboard.
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_A));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_A));
-        waitAndVerifyStatesAndCalls(0, "a", 1, 1, -1, -1);
+        waitAndVerifyUpdateSelection(0, 1, 1, -1, -1);
 
         // Type 'enter' key.
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_ENTER));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_ENTER));
-        waitAndVerifyStatesAndCalls(1, "a\n\n", 2, 2, -1, -1);
+        waitAndVerifyUpdateSelection(1, 2, 2, -1, -1);
+        assertTextsAroundCursor("a\n", "", "\n");
 
         // Type 'b'.
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_B));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_B));
-        waitAndVerifyStatesAndCalls(2, "a\nb", 3, 3, -1, -1);
+        waitAndVerifyUpdateSelection(2, 3, 3, -1, -1);
+        assertTextsAroundCursor("a\nb", "", "");
     }
 
     @SmallTest
@@ -784,72 +812,70 @@ public class ImeTest extends ContentShellTestBase {
         // h
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_H));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_H));
-        assertEquals("h", mConnection.getTextBeforeCursor(9, 0));
-        waitAndVerifyStatesAndCalls(0, "h", 1, 1, -1, -1);
+        assertEquals("h", getTextBeforeCursor(9, 0));
+        waitAndVerifyUpdateSelection(0, 1, 1, -1, -1);
 
         // ALT-i  (circumflex accent key on virtual keyboard)
         dispatchKeyEvent(new KeyEvent(
                 0, 0, KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_I, 0, KeyEvent.META_ALT_ON));
-        assertUpdateStateCall(1000);
-        assertEquals("hˆ", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("hˆ", getTextBeforeCursor(9, 0));
         dispatchKeyEvent(new KeyEvent(
                 0, 0, KeyEvent.ACTION_UP, KeyEvent.KEYCODE_I, 0, KeyEvent.META_ALT_ON));
-        assertEquals("hˆ", mConnection.getTextBeforeCursor(9, 0));
-        waitAndVerifyStatesAndCalls(1, "hˆ", 2, 2, 1, 2);
+        assertEquals("hˆ", getTextBeforeCursor(9, 0));
+        waitAndVerifyUpdateSelection(1, 2, 2, 1, 2);
 
         // o
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_O));
-        assertUpdateStateCall(1000);
-        assertEquals("hô", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("hô", getTextBeforeCursor(9, 0));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_O));
-        assertEquals("hô", mConnection.getTextBeforeCursor(9, 0));
-        waitAndVerifyStatesAndCalls(2, "hô", 2, 2, -1, -1);
+        assertEquals("hô", getTextBeforeCursor(9, 0));
+        waitAndVerifyUpdateSelection(2, 2, 2, -1, -1);
 
         // ALT-i
         dispatchKeyEvent(new KeyEvent(
                 0, 0, KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_I, 0, KeyEvent.META_ALT_ON));
-        assertUpdateStateCall(1000);
         dispatchKeyEvent(new KeyEvent(
                 0, 0, KeyEvent.ACTION_UP, KeyEvent.KEYCODE_I, 0, KeyEvent.META_ALT_ON));
-        assertEquals("hôˆ", mConnection.getTextBeforeCursor(9, 0));
-        waitAndVerifyStatesAndCalls(3, "hôˆ", 3, 3, 2, 3);
+        assertEquals("hôˆ", getTextBeforeCursor(9, 0));
+        waitAndVerifyUpdateSelection(3, 3, 3, 2, 3);
 
         // ALT-i again should have no effect
         dispatchKeyEvent(new KeyEvent(
                 0, 0, KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_I, 0, KeyEvent.META_ALT_ON));
-        assertUpdateStateCall(1000);
         dispatchKeyEvent(new KeyEvent(
                 0, 0, KeyEvent.ACTION_UP, KeyEvent.KEYCODE_I, 0, KeyEvent.META_ALT_ON));
-        assertEquals("hôˆ", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("hôˆ", getTextBeforeCursor(9, 0));
 
         // b (cannot be accented, should just appear after)
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_B));
-        assertUpdateStateCall(1000);
-        assertEquals("hôˆb", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("hôˆb", getTextBeforeCursor(9, 0));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_B));
-        assertEquals("hôˆb", mConnection.getTextBeforeCursor(9, 0));
-        // A transitional state due to finishComposingText.
-        waitAndVerifyStates(4, "hôˆ", 3, 3, -1, -1);
-        waitAndVerifyStatesAndCalls(5, "hôˆb", 4, 4, -1, -1);
+        assertEquals("hôˆb", getTextBeforeCursor(9, 0));
+        int index = 4;
+        if (usingReplicaInputConnection()) {
+            // A transitional state due to finishComposingText.
+            waitAndVerifyUpdateSelection(index++, 3, 3, -1, -1);
+        }
+        waitAndVerifyUpdateSelection(index++, 4, 4, -1, -1);
 
         // ALT-i
         dispatchKeyEvent(new KeyEvent(
                 0, 0, KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_I, 0, KeyEvent.META_ALT_ON));
-        assertUpdateStateCall(1000);
         dispatchKeyEvent(new KeyEvent(
                 0, 0, KeyEvent.ACTION_UP, KeyEvent.KEYCODE_I, 0, KeyEvent.META_ALT_ON));
-        assertEquals("hôˆbˆ", mConnection.getTextBeforeCursor(9, 0));
-        waitAndVerifyStatesAndCalls(6, "hôˆbˆ", 5, 5, 4, 5);
+        assertEquals("hôˆbˆ", getTextBeforeCursor(9, 0));
+        waitAndVerifyUpdateSelection(index++, 5, 5, 4, 5);
 
         // Backspace
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_DEL));
-        assertUpdateStateCall(1000);
-        assertEquals("hôˆb", mConnection.getTextBeforeCursor(9, 0));
+        assertEquals("hôˆb", getTextBeforeCursor(9, 0));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_DEL));
-        assertEquals("hôˆb", mConnection.getTextBeforeCursor(9, 0));
-        // A transitional state due to finishComposingText in deleteSurroundingTextImpl.
-        waitAndVerifyStates(7, "hôˆbˆ", 5, 5, -1, -1);
-        waitAndVerifyStatesAndCalls(8, "hôˆb", 4, 4, -1, -1);
+        assertEquals("hôˆb", getTextBeforeCursor(9, 0));
+        if (usingReplicaInputConnection()) {
+            // A transitional state due to finishComposingText in deleteSurroundingTextImpl.
+            waitAndVerifyUpdateSelection(index++, 5, 5, -1, -1);
+        }
+        waitAndVerifyUpdateSelection(index++, 4, 4, -1, -1);
     }
 
     @SmallTest
@@ -869,17 +895,16 @@ public class ImeTest extends ContentShellTestBase {
         focusElementAndWaitForStateUpdate("textarea");
 
         commitText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, -1, -1);
+        waitAndVerifyUpdateSelection(0, 5, 5, -1, -1);
 
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_ENTER));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_ENTER));
-        // TODO(aurimas): remove this workaround when crbug.com/278584 is fixed.
-        // The second new line is not a user visible/editable one, it is a side-effect of Blink
-        // using <br> internally. This only happens when \n is at the end.
-        waitAndVerifyStatesAndCalls(1, "hello\n\n", 6, 6, -1, -1);
+        waitAndVerifyUpdateSelection(1, 6, 6, -1, -1);
+        assertTextsAroundCursor("hello\n", "", "\n");
 
         commitText("world", 1);
-        waitAndVerifyStatesAndCalls(2, "hello\nworld", 11, 11, -1, -1);
+        waitAndVerifyUpdateSelection(2, 11, 11, -1, -1);
+        assertTextsAroundCursor("hello\nworld", "", "");
     }
 
     @SmallTest
@@ -888,20 +913,22 @@ public class ImeTest extends ContentShellTestBase {
         focusElementAndWaitForStateUpdate("textarea");
 
         setComposingText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, 0, 5);
+        waitAndVerifyUpdateSelection(0, 5, 5, 0, 5);
+
+        // IME app should call this, otherwise enter key should clear the current composition.
+        finishComposingText();
+        waitAndVerifyUpdateSelection(1, 5, 5, -1, -1);
 
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_ENTER));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_ENTER));
 
-        // TODO(aurimas): remove this workaround when crbug.com/278584 is fixed.
-        // A transitional state due to finishComposingText.
-        waitAndVerifyStates(1, "hello", 5, 5, -1, -1);
         // The second new line is not a user visible/editable one, it is a side-effect of Blink
         // using <br> internally. This only happens when \n is at the end.
-        waitAndVerifyStatesAndCalls(2, "hello\n\n", 6, 6, -1, -1);
+        waitAndVerifyUpdateSelection(2, 6, 6, -1, -1);
 
         commitText("world", 1);
-        waitAndVerifyStatesAndCalls(3, "hello\nworld", 11, 11, -1, -1);
+        waitAndVerifyUpdateSelection(3, 11, 11, -1, -1);
+        assertTextsAroundCursor("hello\nworld", "", "");
     }
 
     @SmallTest
@@ -910,9 +937,9 @@ public class ImeTest extends ContentShellTestBase {
         focusElement("textarea");
 
         // DPAD_CENTER should cause keyboard to appear
-        resetUpdateStateList();
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_DPAD_CENTER));
-        assertUpdateStateCall(1000);
+
+        // TODO(changwan): should really check this.
     }
 
     @SmallTest
@@ -921,21 +948,24 @@ public class ImeTest extends ContentShellTestBase {
         focusElementAndWaitForStateUpdate("textarea");
 
         commitText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, -1, -1);
+        waitAndVerifyUpdateSelection(0, 5, 5, -1, -1);
 
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_DOWN, KeyEvent.KEYCODE_DPAD_LEFT));
         dispatchKeyEvent(new KeyEvent(KeyEvent.ACTION_UP, KeyEvent.KEYCODE_DPAD_LEFT));
 
-        // Ideally getTextBeforeCursor immediately after dispatchKeyEvent should return a correct
-        // value, but we have a stop-gap solution in render_widget_input_handler and it make take
-        // some round trip time until we get the correct value.
-        // TODO(changwan): Change the test not to wait when we fix it.
-        waitUntilGetCharacterBeforeCursorBecomes("l");
+        if (usingReplicaInputConnection()) {
+            // Ideally getTextBeforeCursor immediately after dispatchKeyEvent should return a
+            // correct value, but we have a stop-gap solution in render_widget_input_handler and it
+            // make take some round trip time until we get the correct value.
+            waitUntilGetCharacterBeforeCursorBecomes("l");
+        } else {
+            assertTextsAroundCursor("hell", "", "o");
+        }
     }
 
     private void waitUntilGetCharacterBeforeCursorBecomes(final String expectedText)
             throws InterruptedException {
-        CriteriaHelper.pollForUIThreadCriteria(new Criteria() {
+        pollForCriteriaOnThread(new Criteria() {
             @Override
             public boolean isSatisfied() {
                 String actualText = (String) mConnection.getTextBeforeCursor(1, 0);
@@ -945,17 +975,45 @@ public class ImeTest extends ContentShellTestBase {
         });
     }
 
+    private void pollForCriteriaOnThread(final Criteria criteria) throws InterruptedException {
+        final Callable<Boolean> callable = new Callable<Boolean>() {
+            @Override
+            public Boolean call() throws Exception {
+                return criteria.isSatisfied();
+            }
+        };
+        CriteriaHelper.pollForCriteria(new Criteria() {
+            @Override
+            public boolean isSatisfied() {
+                try {
+                    return runBlockingOnImeThread(callable);
+                } catch (Exception e) {
+                    e.printStackTrace();
+                    fail();
+                    return false;
+                }
+            }
+
+            @Override
+            public String getFailureReason() {
+                return criteria.getFailureReason();
+            }
+        });
+    }
+
     @SmallTest
     @Feature({"TextInput"})
     public void testPastePopupShowAndHide() throws Throwable {
         commitText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, -1, -1);
+        waitAndVerifyUpdateSelection(0, 5, 5, -1, -1);
 
         selectAll();
-        waitAndVerifyStatesAndCalls(1, "hello", 0, 5, -1, -1);
+        waitAndVerifyUpdateSelection(1, 0, 5, -1, -1);
+        assertTextsAroundCursor("", "hello", "");
 
         cut();
-        waitAndVerifyStatesAndCalls(2, "", 0, 0, -1, -1);
+        waitAndVerifyUpdateSelection(2, 0, 0, -1, -1);
+        assertTextsAroundCursor("", "", "");
 
         DOMUtils.longPressNode(this, mContentViewCore, "input_text");
         CriteriaHelper.pollForUIThreadCriteria(new Criteria() {
@@ -982,7 +1040,7 @@ public class ImeTest extends ContentShellTestBase {
     @Feature({"TextInput"})
     public void testSelectionClearedOnKeyEvent() throws Throwable {
         commitText("hello", 1);
-        waitAndVerifyStatesAndCalls(0, "hello", 5, 5, -1, -1);
+        waitAndVerifyUpdateSelection(0, 5, 5, -1, -1);
 
         DOMUtils.clickNode(this, mContentViewCore, "input_text");
         assertWaitForKeyboardStatus(true);
@@ -1012,28 +1070,45 @@ public class ImeTest extends ContentShellTestBase {
     @Feature({"TextInput"})
     public void testRestartInputWhileComposingText() throws Throwable {
         setComposingText("abc", 1);
-        waitAndVerifyStatesAndCalls(0, "abc", 3, 3, 0, 3);
-        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
-            @Override
-            public void run() {
-                mImeAdapter.restartInput();
-            }
-        });
+        waitAndVerifyUpdateSelection(0, 3, 3, 0, 3);
+        restartInput();
         // We don't do anything when input gets restarted. But we depend on Android's
         // InputMethodManager and/or input methods to call finishComposingText() in setting
         // current input connection as active or finishing the current input connection.
-        assertNoFurtherStateUpdate(1);
+        Thread.sleep(1000);
+        assertEquals(1, mInputMethodManagerWrapper.getUpdateSelectionList().size());
+    }
+
+    @MediumTest
+    @Feature({"TextInput"})
+    public void testRestartInputKeepsTextAndCursor() throws Exception {
+        commitText("ab", 2);
+        restartInput();
+        assertEquals("ab", getTextBeforeCursor(10, 0));
     }
 
     private void performGo(TestCallbackHelperContainer testCallbackHelperContainer)
             throws Throwable {
-        final AdapterInputConnection inputConnection = mConnection;
+        final InputConnection inputConnection = mConnection;
+        final Callable<Void> callable = new Callable<Void>() {
+            @Override
+            public Void call() throws Exception {
+                inputConnection.performEditorAction(EditorInfo.IME_ACTION_GO);
+                return null;
+            }
+        };
+
         handleBlockingCallbackAction(
                 testCallbackHelperContainer.getOnPageFinishedHelper(),
                 new Runnable() {
                     @Override
                     public void run() {
-                        inputConnection.performEditorAction(EditorInfo.IME_ACTION_GO);
+                        try {
+                            runBlockingOnImeThread(callable);
+                        } catch (Exception e) {
+                            e.printStackTrace();
+                            fail();
+                        }
                     }
                 });
     }
@@ -1044,8 +1119,7 @@ public class ImeTest extends ContentShellTestBase {
             public boolean isSatisfied() {
                 // We do not check the other way around: in some cases we need to keep
                 // input connection even when the last known status is 'hidden'.
-                // See crbug.com/569332 for more details.
-                if (show && getAdapterInputConnection() == null) {
+                if (show && getInputConnection() == null) {
                     updateFailureReason("input connection should not be null.");
                     return false;
                 }
@@ -1065,56 +1139,25 @@ public class ImeTest extends ContentShellTestBase {
         });
     }
 
-    private void waitAndVerifyStates(final int index, String text, final int selectionStart,
+    private void waitAndVerifyUpdateSelection(final int index, final int selectionStart,
             final int selectionEnd, final int compositionStart, final int compositionEnd)
             throws InterruptedException {
-        final List<TestImeState> states = mConnectionFactory.getImeStateList();
+        final List<Pair<Range, Range>> states = mInputMethodManagerWrapper.getUpdateSelectionList();
         CriteriaHelper.pollForUIThreadCriteria(new Criteria() {
             @Override
             public boolean isSatisfied() {
                 return states.size() > index;
             }
         });
-        states.get(index).assertEqualState(
-                text, selectionStart, selectionEnd, compositionStart, compositionEnd);
+        Pair<Range, Range> selection = states.get(index);
+        assertEquals(selectionStart, selection.first.start());
+        assertEquals(selectionEnd, selection.first.end());
+        assertEquals(compositionStart, selection.second.start());
+        assertEquals(compositionEnd, selection.second.end());
     }
 
-    private void waitAndVerifyStatesAndCalls(final int index, String text, final int selectionStart,
-            final int selectionEnd, final int compositionStart, final int compositionEnd)
-            throws InterruptedException {
-        waitAndVerifyStates(
-                index, text, selectionStart, selectionEnd, compositionStart, compositionEnd);
-
-        // Wait and verify calls to InputMethodManager.
-        final Range selection = new Range(selectionStart, selectionEnd);
-        final Range composition = new Range(compositionStart, compositionEnd);
-        CriteriaHelper.pollForUIThreadCriteria(new Criteria() {
-            @Override
-            public boolean isSatisfied() {
-                updateFailureReason(
-                        "Actual selection was: " + mInputMethodManagerWrapper.getSelection()
-                        + ", and actual composition was: "
-                        + mInputMethodManagerWrapper.getComposition());
-                return mInputMethodManagerWrapper.getSelection().equals(selection)
-                        && mInputMethodManagerWrapper.getComposition().equals(composition);
-            }
-        });
-    }
-
-    private void resetUpdateStateList() {
-        mConnectionFactory.getImeStateList().clear();
-    }
-
-    private void assertUpdateStateCall(int maxms) throws Exception {
-        while (mConnectionFactory.getImeStateList().size() == 0 && maxms > 0) {
-            try {
-                Thread.sleep(50);
-            } catch (Exception e) {
-                // Not really a problem since we're just going to sleep again.
-            }
-            maxms -= 50;
-        }
-        assertTrue(mConnectionFactory.getImeStateList().size() > 0);
+    private void resetUpdateSelectionList() {
+        mInputMethodManagerWrapper.getUpdateSelectionList().clear();
     }
 
     private void assertClipboardContents(final Activity activity, final String expectedContents)
@@ -1135,8 +1178,28 @@ public class ImeTest extends ContentShellTestBase {
         return mContentViewCore.getImeAdapterForTest();
     }
 
-    private AdapterInputConnection getAdapterInputConnection() {
-        return mContentViewCore.getImeAdapterForTest().getInputConnectionForTest();
+    private ChromiumBaseInputConnection getInputConnection() {
+        try {
+            return ThreadUtils.runOnUiThreadBlocking(new Callable<ChromiumBaseInputConnection>() {
+                @Override
+                public ChromiumBaseInputConnection call() {
+                    return mContentViewCore.getImeAdapterForTest().getInputConnectionForTest();
+                }
+            });
+        } catch (ExecutionException e) {
+            e.printStackTrace();
+            fail();
+            return null;
+        }
+    }
+
+    private void restartInput() {
+        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
+            @Override
+            public void run() {
+                mImeAdapter.restartInput();
+            }
+        });
     }
 
     private void copy() {
@@ -1189,62 +1252,98 @@ public class ImeTest extends ContentShellTestBase {
         });
     }
 
-    private void commitText(final CharSequence text, final int newCursorPosition) {
-        final AdapterInputConnection connection = mConnection;
-        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
+    private <T> T runBlockingOnImeThread(Callable<T> c) throws Exception {
+        return ImeTestUtils.runBlockingOnHandler(mConnectionFactory.getHandler(), c);
+    }
+
+    private boolean commitText(final CharSequence text, final int newCursorPosition)
+            throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<Boolean>() {
             @Override
-            public void run() {
-                connection.commitText(text, newCursorPosition);
+            public Boolean call() {
+                return connection.commitText(text, newCursorPosition);
             }
         });
     }
 
-    private void setSelection(final int start, final int end) {
-        final AdapterInputConnection connection = mConnection;
-        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
+    private boolean setSelection(final int start, final int end) throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<Boolean>() {
             @Override
-            public void run() {
-                connection.setSelection(start, end);
+            public Boolean call() {
+                return connection.setSelection(start, end);
             }
         });
     }
 
-    private void setComposingRegion(final int start, final int end) {
-        final AdapterInputConnection connection = mConnection;
-        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
+    private boolean setComposingRegion(final int start, final int end) throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<Boolean>() {
             @Override
-            public void run() {
-                connection.setComposingRegion(start, end);
+            public Boolean call() {
+                return connection.setComposingRegion(start, end);
             }
         });
     }
 
-    private void setComposingText(final CharSequence text, final int newCursorPosition) {
-        final AdapterInputConnection connection = mConnection;
-        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
+    private boolean setComposingText(final CharSequence text, final int newCursorPosition)
+            throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<Boolean>() {
             @Override
-            public void run() {
-                connection.setComposingText(text, newCursorPosition);
+            public Boolean call() {
+                return connection.setComposingText(text, newCursorPosition);
             }
         });
     }
 
-    private void finishComposingText() {
-        final AdapterInputConnection connection = mConnection;
-        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
+    private boolean finishComposingText() throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<Boolean>() {
             @Override
-            public void run() {
-                connection.finishComposingText();
+            public Boolean call() {
+                return connection.finishComposingText();
             }
         });
     }
 
-    private void deleteSurroundingText(final int before, final int after) {
-        final AdapterInputConnection connection = mConnection;
-        ThreadUtils.runOnUiThreadBlocking(new Runnable() {
+    private boolean deleteSurroundingText(final int before, final int after) throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<Boolean>() {
             @Override
-            public void run() {
-                connection.deleteSurroundingText(before, after);
+            public Boolean call() {
+                return connection.deleteSurroundingText(before, after);
+            }
+        });
+    }
+
+    private CharSequence getTextBeforeCursor(final int length, final int flags) throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<CharSequence>() {
+            @Override
+            public CharSequence call() {
+                return connection.getTextBeforeCursor(length, flags);
+            }
+        });
+    }
+
+    private CharSequence getSelectedText(final int flags) throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<CharSequence>() {
+            @Override
+            public CharSequence call() {
+                return connection.getSelectedText(flags);
+            }
+        });
+    }
+
+    private CharSequence getTextAfterCursor(final int length, final int flags) throws Exception {
+        final ChromiumBaseInputConnection connection = mConnection;
+        return runBlockingOnImeThread(new Callable<CharSequence>() {
+            @Override
+            public CharSequence call() {
+                return connection.getTextAfterCursor(length, flags);
             }
         });
     }
@@ -1265,8 +1364,8 @@ public class ImeTest extends ContentShellTestBase {
     private void focusElementAndWaitForStateUpdate(String id)
             throws InterruptedException, TimeoutException {
         focusElement(id);
-        waitAndVerifyStatesAndCalls(0, "", 0, 0, -1, -1);
-        resetUpdateStateList();
+        waitAndVerifyUpdateSelection(0, 0, 0, -1, -1);
+        resetUpdateSelectionList();
     }
 
     private void focusElement(final String id) throws InterruptedException, TimeoutException {
@@ -1288,24 +1387,36 @@ public class ImeTest extends ContentShellTestBase {
             }
         });
         // When we focus another element, the connection may be recreated.
-        mConnection = (TestAdapterInputConnection) getAdapterInputConnection();
+        mConnection = getInputConnection();
+    }
+
+    private boolean usingReplicaInputConnection() {
+        return mConnectionFactory.getHandler().getLooper() == Looper.getMainLooper();
     }
 
-    private static class TestAdapterInputConnectionFactory extends
-            ImeAdapter.AdapterInputConnectionFactory {
-        private final List<TestImeState> mImeStateList = new ArrayList<>();
+    private static class TestInputConnectionFactory implements ChromiumBaseInputConnection.Factory {
+        private final ChromiumBaseInputConnection.Factory mFactory;
+
         private final List<Integer> mTextInputTypeList = new ArrayList<>();
+        private EditorInfo mOutAttrs;
+
+        public TestInputConnectionFactory(ChromiumBaseInputConnection.Factory factory) {
+            mFactory = factory;
+        }
 
         @Override
-        public AdapterInputConnection get(View view, ImeAdapter imeAdapter, int initialSelStart,
-                int initialSelEnd, EditorInfo outAttrs) {
-            mTextInputTypeList.add(imeAdapter.getTextInputType());
-            return new TestAdapterInputConnection(
-                    mImeStateList, view, imeAdapter, initialSelStart, initialSelEnd, outAttrs);
+        public ChromiumBaseInputConnection initializeAndGet(View view, ImeAdapter imeAdapter,
+                int inputType, int inputFlags, int selectionStart, int selectionEnd,
+                EditorInfo outAttrs) {
+            mTextInputTypeList.add(inputType);
+            mOutAttrs = outAttrs;
+            return mFactory.initializeAndGet(view, imeAdapter, inputType, inputFlags,
+                    selectionStart, selectionEnd, outAttrs);
         }
 
-        public List<TestImeState> getImeStateList() {
-            return mImeStateList;
+        @Override
+        public Handler getHandler() {
+            return mFactory.getHandler();
         }
 
         public Integer[] getTextInputTypeHistory() {
@@ -1317,51 +1428,9 @@ public class ImeTest extends ContentShellTestBase {
         public void clearTextInputTypeHistory() {
             mTextInputTypeList.clear();
         }
-    }
-
-    private static class TestAdapterInputConnection extends AdapterInputConnection {
-        private final List<TestImeState> mImeStateList;
-
-        public TestAdapterInputConnection(List<TestImeState> imeStateList, View view,
-                ImeAdapter imeAdapter, int initialSelStart, int initialSelEnd,
-                EditorInfo outAttrs) {
-            super(view, imeAdapter, initialSelStart, initialSelEnd, outAttrs);
-            mImeStateList = imeStateList;
-        }
-
-        @Override
-        public void updateState(String text, int selectionStart, int selectionEnd,
-                int compositionStart, int compositionEnd, boolean requiredAck) {
-            mImeStateList.add(new TestImeState(
-                    text, selectionStart, selectionEnd, compositionStart, compositionEnd));
-            super.updateState(text, selectionStart, selectionEnd, compositionStart,
-                    compositionEnd, requiredAck);
-        }
-    }
-
-    private static class TestImeState {
-        private final String mText;
-        private final int mSelectionStart;
-        private final int mSelectionEnd;
-        private final int mCompositionStart;
-        private final int mCompositionEnd;
-
-        public TestImeState(String text, int selectionStart, int selectionEnd,
-                int compositionStart, int compositionEnd) {
-            mText = text;
-            mSelectionStart = selectionStart;
-            mSelectionEnd = selectionEnd;
-            mCompositionStart = compositionStart;
-            mCompositionEnd = compositionEnd;
-        }
 
-        public void assertEqualState(String text, int selectionStart, int selectionEnd,
-                int compositionStart, int compositionEnd) {
-            assertEquals("Text did not match", text, mText);
-            assertEquals("Selection start did not match", selectionStart, mSelectionStart);
-            assertEquals("Selection end did not match", selectionEnd, mSelectionEnd);
-            assertEquals("Composition start did not match", compositionStart, mCompositionStart);
-            assertEquals("Composition end did not match", compositionEnd, mCompositionEnd);
+        public EditorInfo getOutAttrs() {
+            return mOutAttrs;
         }
     }
 }