2 files changed, 58 insertions, 9 deletions

diff --git a/tests/src/com/android/inputmethod/latin/BinaryDictionaryTests.java b/tests/src/com/android/inputmethod/latin/BinaryDictionaryTests.java
index 7ed3ee180..826c0f7b2 100644
--- a/tests/src/com/android/inputmethod/latin/BinaryDictionaryTests.java
+++ b/tests/src/com/android/inputmethod/latin/BinaryDictionaryTests.java
@@ -27,6 +27,7 @@ import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Random;
@@ -625,4 +626,57 @@ public class BinaryDictionaryTests extends AndroidTestCase {
 
         dictFile.delete();
     }
+
+    public void testUnigramAndBigramCount() {
+        final int flashWithGCIterationCount = 10;
+        final int codePointSetSize = 50;
+        final int unigramCountPerIteration = 1000;
+        final int bigramCountPerIteration = 2000;
+        final int seed = 1123581321;
+
+        final Random random = new Random(seed);
+
+        File dictFile = null;
+        try {
+            dictFile = createEmptyDictionaryAndGetFile("TestBinaryDictionary");
+        } catch (IOException e) {
+            fail("IOException while writing an initial dictionary : " + e);
+        }
+
+        final ArrayList<String> words = new ArrayList<String>();
+        final HashSet<Pair<String, String>> bigrams = new HashSet<Pair<String, String>>();
+        final int[] codePointSet = CodePointUtils.generateCodePointSet(codePointSetSize, random);
+
+        BinaryDictionary binaryDictionary;
+        for (int i = 0; i < flashWithGCIterationCount; i++) {
+            binaryDictionary = new BinaryDictionary(dictFile.getAbsolutePath(),
+                    0 /* offset */, dictFile.length(), true /* useFullEditDistance */,
+                    Locale.getDefault(), TEST_LOCALE, true /* isUpdatable */);
+            for (int j = 0; j < unigramCountPerIteration; j++) {
+                final String word = CodePointUtils.generateWord(random, codePointSet);
+                words.add(word);
+                final int unigramProbability = random.nextInt(0xFF);
+                binaryDictionary.addUnigramWord(word, unigramProbability);
+            }
+            for (int j = 0; j < bigramCountPerIteration; j++) {
+                final String word0 = words.get(random.nextInt(words.size()));
+                final String word1 = words.get(random.nextInt(words.size()));
+                bigrams.add(new Pair<String, String>(word0, word1));
+                final int bigramProbability = random.nextInt(0xF);
+                binaryDictionary.addBigramWords(word0, word1, bigramProbability);
+            }
+            assertEquals(new HashSet<String>(words).size(), Integer.parseInt(
+                    binaryDictionary.getPropertyForTests(BinaryDictionary.UNIGRAM_COUNT_QUERY)));
+            assertEquals(new HashSet<Pair<String, String>>(bigrams).size(), Integer.parseInt(
+                    binaryDictionary.getPropertyForTests(BinaryDictionary.BIGRAM_COUNT_QUERY)));
+            binaryDictionary.flushWithGC();
+            assertEquals(new HashSet<String>(words).size(), Integer.parseInt(
+                    binaryDictionary.getPropertyForTests(BinaryDictionary.UNIGRAM_COUNT_QUERY)));
+            assertEquals(new HashSet<Pair<String, String>>(bigrams).size(), Integer.parseInt(
+                    binaryDictionary.getPropertyForTests(BinaryDictionary.BIGRAM_COUNT_QUERY)));
+            binaryDictionary.close();
+        }
+
+        dictFile.delete();
+    }
 }
diff --git a/tests/src/com/android/inputmethod/latin/personalization/UserHistoryDictionaryTests.java b/tests/src/com/android/inputmethod/latin/personalization/UserHistoryDictionaryTests.java
index cf649c489..ddc9546c5 100644
--- a/tests/src/com/android/inputmethod/latin/personalization/UserHistoryDictionaryTests.java
+++ b/tests/src/com/android/inputmethod/latin/personalization/UserHistoryDictionaryTests.java
@@ -84,29 +84,24 @@ public class UserHistoryDictionaryTests extends AndroidTestCase {
     }
 
     /**
-     * @param checksContents if true, checks whether written words are actually in the dictionary
+     * @param checkContents if true, checks whether written words are actually in the dictionary
      * or not.
      */
     private void addAndWriteRandomWords(final String testFilenameSuffix, final int numberOfWords,
-            final Random random, final boolean checksContents) {
+            final Random random, final boolean checkContents) {
         final List<String> words = generateWords(numberOfWords, random);
         final UserHistoryDictionary dict =
                 PersonalizationHelper.getUserHistoryDictionary(getContext(),
                         testFilenameSuffix /* locale */, mPrefs);
         // Add random words to the user history dictionary.
         addToDict(dict, words);
-        if (checksContents) {
+        if (checkContents) {
             try {
                 Thread.sleep(TimeUnit.MILLISECONDS.convert(5L, TimeUnit.SECONDS));
             } catch (InterruptedException e) {
             }
-            // Limit word count to check when using a Java on memory dictionary.
-            final int wordCountToCheck =
-                    ExpandableBinaryDictionary.ENABLE_BINARY_DICTIONARY_DYNAMIC_UPDATE ?
-                            numberOfWords : 10;
-            for (int i = 0; i < wordCountToCheck; ++i) {
+            for (int i = 0; i < numberOfWords; ++i) {
                 final String word = words.get(i);
-                // This may fail as long as we use tryLock on inserting the bigram words
                 assertTrue(dict.isInDictionaryForTests(word));
             }
         }