Add test only methods to test personalized dicts.

Bug: 11122480
Change-Id: I79a4a73ff09a2d34d4ca094d705e901fb4bc855c

3 files changed, 33 insertions, 6 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();
diff --git a/tests/src/com/android/inputmethod/latin/InputTestsBase.java b/tests/src/com/android/inputmethod/latin/InputTestsBase.java
index 4b157e700..690c559e8 100644
--- a/tests/src/com/android/inputmethod/latin/InputTestsBase.java
+++ b/tests/src/com/android/inputmethod/latin/InputTestsBase.java
@@ -275,9 +275,9 @@ public class InputTestsBase extends ServiceTestCase<LatinIMEForTests> {
         }
     }
 
-    protected void waitForDictionaryToBeLoaded() {
+    protected void waitForDictionariesToBeLoaded() {
         try {
-            mLatinIME.waitForMainDictionary(
+            mLatinIME.waitForLoadingDictionaries(
                     TIMEOUT_TO_WAIT_FOR_LOADING_MAIN_DICTIONARY_IN_SECONDS, TimeUnit.SECONDS);
         } catch (InterruptedException e) {
             Log.e(TAG, "Interrupted during waiting for loading main dictionary.", e);
@@ -286,7 +286,7 @@ public class InputTestsBase extends ServiceTestCase<LatinIMEForTests> {
 
     protected void changeLanguage(final String locale) {
         changeLanguageWithoutWait(locale);
-        waitForDictionaryToBeLoaded();
+        waitForDictionariesToBeLoaded();
     }
 
     protected void changeLanguageWithoutWait(final String locale) {
@@ -314,6 +314,7 @@ public class InputTestsBase extends ServiceTestCase<LatinIMEForTests> {
         mLatinIME.loadKeyboard();
         runMessages();
         mKeyboard = mLatinIME.mKeyboardSwitcher.getKeyboard();
+        mLatinIME.clearPersonalizedDictionariesForTest();
     }
 
     protected void changeKeyboardLocaleAndDictLocale(final String keyboardLocale,
@@ -322,7 +323,7 @@ public class InputTestsBase extends ServiceTestCase<LatinIMEForTests> {
         if (!keyboardLocale.equals(dictLocale)) {
             mLatinIME.replaceDictionariesForTest(LocaleUtils.constructLocaleFromString(dictLocale));
         }
-        waitForDictionaryToBeLoaded();
+        waitForDictionariesToBeLoaded();
     }
 
     protected void pickSuggestionManually(final int index, final String suggestion) {