diff --git a/java/src/com/android/inputmethod/latin/personalization/DynamicPredictionDictionaryBase.java b/java/src/com/android/inputmethod/latin/personalization/DynamicPredictionDictionaryBase.java
index b2e979d..a8600c0 100644
--- a/java/src/com/android/inputmethod/latin/personalization/DynamicPredictionDictionaryBase.java
+++ b/java/src/com/android/inputmethod/latin/personalization/DynamicPredictionDictionaryBase.java
@@ -404,6 +404,7 @@
     public void registerUpdateSession(PersonalizationDictionaryUpdateSession session) {
         session.setPredictionDictionary(mLocale, this);
         mSessions.add(session);
+        session.onDictionaryReady();
     }
 
     public void unRegisterUpdateSession(PersonalizationDictionaryUpdateSession session) {
diff --git a/tests/src/com/android/inputmethod/latin/InputTestsBase.java b/tests/src/com/android/inputmethod/latin/InputTestsBase.java
index eb4f706..ead134b 100644
--- a/tests/src/com/android/inputmethod/latin/InputTestsBase.java
+++ b/tests/src/com/android/inputmethod/latin/InputTestsBase.java
@@ -204,17 +204,16 @@
         // view and only delegates to the parts of the code that care. So we don't include them here
         // to keep these tests as pinpoint as possible and avoid bringing it too many dependencies,
         // but keep them in mind if something breaks. Commenting them out as is should work.
-        //mLatinIME.onPressKey(codePoint);
-        for (final Key key : mKeyboard.mKeys) {
-            if (key.mCode == codePoint) {
-                final int x = key.mX + key.mWidth / 2;
-                final int y = key.mY + key.mHeight / 2;
-                mLatinIME.onCodeInput(codePoint, x, y);
-                return;
-            }
+        //mLatinIME.onPressKey(codePoint, 0 /* repeatCount */, true /* isSinglePointer */);
+        final Key key = mKeyboard.getKey(codePoint);
+        if (key != null) {
+            final int x = key.mX + key.mWidth / 2;
+            final int y = key.mY + key.mHeight / 2;
+            mLatinIME.onCodeInput(codePoint, x, y);
+            return;
         }
         mLatinIME.onCodeInput(codePoint, Constants.NOT_A_COORDINATE, Constants.NOT_A_COORDINATE);
-        //mLatinIME.onReleaseKey(codePoint, false);
+        //mLatinIME.onReleaseKey(codePoint, false /* withSliding */);
     }
 
     protected void type(final String stringToType) {