2 files changed, 28 insertions, 2 deletions
diff --git a/java/src/com/android/inputmethod/latin/DictionaryFacilitatorForSuggest.java b/java/src/com/android/inputmethod/latin/DictionaryFacilitatorForSuggest.java
index 138a626a0..43d4ba421 100644
--- a/java/src/com/android/inputmethod/latin/DictionaryFacilitatorForSuggest.java
+++ b/java/src/com/android/inputmethod/latin/DictionaryFacilitatorForSuggest.java
@@ -237,6 +237,24 @@ public class DictionaryFacilitatorForSuggest {
         mLatchForWaitingLoadingMainDictionary.await(timeout, unit);
     }
 
+    @UsedForTesting
+    public void waitForLoadingDictionariesForTesting(final long timeout, final TimeUnit unit)
+            throws InterruptedException {
+        waitForLoadingMainDictionary(timeout, unit);
+        if (mContactsDictionary != null) {
+            mContactsDictionary.waitAllTasksForTests();
+        }
+        if (mUserDictionary != null) {
+            mUserDictionary.waitAllTasksForTests();
+        }
+        if (mUserHistoryDictionary != null) {
+            mUserHistoryDictionary.waitAllTasksForTests();
+        }
+        if (mPersonalizationDictionary != null) {
+            mPersonalizationDictionary.waitAllTasksForTests();
+        }
+    }
+
     private void setMainDictionary(final Dictionary mainDictionary) {
         mMainDictionary = mainDictionary;
         addOrReplaceDictionary(Dictionary.TYPE_MAIN, mainDictionary);
diff --git a/java/src/com/android/inputmethod/latin/LatinIME.java b/java/src/com/android/inputmethod/latin/LatinIME.java
index 346ba8523..47a3e9469 100644
--- a/java/src/com/android/inputmethod/latin/LatinIME.java
+++ b/java/src/com/android/inputmethod/latin/LatinIME.java
@@ -1717,9 +1717,10 @@ public class LatinIME extends InputMethodService implements KeyboardActionListen
 
     // DO NOT USE THIS for any other purpose than testing. This is information private to LatinIME.
     @UsedForTesting
-    /* package for test */ void waitForMainDictionary(final long timeout, final TimeUnit unit)
+    /* package for test */ void waitForLoadingDictionaries(final long timeout, final TimeUnit unit)
             throws InterruptedException {
-        mInputLogic.mSuggest.mDictionaryFacilitator.waitForLoadingMainDictionary(timeout, unit);
+        mInputLogic.mSuggest.mDictionaryFacilitator.waitForLoadingDictionariesForTesting(
+                timeout, unit);
     }
 
     // DO NOT USE THIS for any other purpose than testing. This can break the keyboard badly.
@@ -1733,6 +1734,13 @@ public class LatinIME extends InputMethodService implements KeyboardActionListen
         resetSuggest(new Suggest(locale, dictionaryFacilitator));
     }
 
+    // DO NOT USE THIS for any other purpose than testing.
+    @UsedForTesting
+    /* package for test */ void clearPersonalizedDictionariesForTest() {
+        mInputLogic.mSuggest.mDictionaryFacilitator.clearUserHistoryDictionary();
+        mInputLogic.mSuggest.mDictionaryFacilitator.clearPersonalizationDictionary();
+    }
+
     public void dumpDictionaryForDebug(final String dictName) {
         if (mInputLogic.mSuggest == null) {
             initSuggest();