diff --git a/tests/src/com/android/inputmethod/latin/BinaryDictionaryTests.java b/tests/src/com/android/inputmethod/latin/BinaryDictionaryTests.java
index d9c35c4..4d44135 100644
--- a/tests/src/com/android/inputmethod/latin/BinaryDictionaryTests.java
+++ b/tests/src/com/android/inputmethod/latin/BinaryDictionaryTests.java
@@ -948,10 +948,13 @@
             final HashSet<String> bigramWord1s = bigrams.get(word0);
             final WordProperty wordProperty = binaryDictionary.getWordProperty(word0);
             assertEquals(bigramWord1s.size(), wordProperty.mBigrams.size());
+            final int unigramProbability = wordProperty.getProbability();
             for (int j = 0; j < wordProperty.mBigrams.size(); j++) {
                 final String word1 = wordProperty.mBigrams.get(j).mWord;
                 assertTrue(bigramWord1s.contains(word1));
-                final int probability = wordProperty.mBigrams.get(j).getProbability();
+                final int bigramProbability = wordProperty.mBigrams.get(j).getProbability();
+                final int probability = binaryDictionary.calculateProbability(
+                        unigramProbability, bigramProbability);
                 assertEquals((int)bigramProbabilities.get(new Pair<String, String>(word0, word1)),
                         probability);
                 assertEquals(wordProperty.mBigrams.get(j).getProbability(), probability);
@@ -1037,11 +1040,14 @@
             assertEquals((int)wordProbabilitiesToCheckLater.get(word0),
                     wordProperty.mProbabilityInfo.mProbability);
             wordSet.remove(word0);
+            final int unigramProbability = wordProperty.getProbability();
             final HashSet<String> bigramWord1s = bigrams.get(word0);
             for (int j = 0; j < wordProperty.mBigrams.size(); j++) {
                 final String word1 = wordProperty.mBigrams.get(j).mWord;
                 assertTrue(bigramWord1s.contains(word1));
-                final int probability = wordProperty.mBigrams.get(j).getProbability();
+                final int bigramProbability = wordProperty.mBigrams.get(j).getProbability();
+                final int probability = binaryDictionary.calculateProbability(
+                        unigramProbability, bigramProbability);
                 final Pair<String, String> bigram = new Pair<String, String>(word0, word1);
                 assertEquals((int)bigramProbabilitiesToCheckLater.get(bigram), probability);
                 bigramSet.remove(bigram);