1 files changed, 13 insertions, 0 deletions
diff --git a/java/src/com/android/inputmethod/research/MainLogBuffer.java b/java/src/com/android/inputmethod/research/MainLogBuffer.java
index 45b83dd76..9e77a1a38 100644
--- a/java/src/com/android/inputmethod/research/MainLogBuffer.java
+++ b/java/src/com/android/inputmethod/research/MainLogBuffer.java
@@ -18,6 +18,7 @@ package com.android.inputmethod.research;
 
 import android.util.Log;
 
+import com.android.inputmethod.annotations.UsedForTesting;
 import com.android.inputmethod.latin.Dictionary;
 import com.android.inputmethod.latin.Suggest;
 import com.android.inputmethod.latin.define.ProductionFlag;
@@ -64,7 +65,11 @@ public abstract class MainLogBuffer extends FixedLogBuffer {
     // The size of the n-grams logged.  E.g. N_GRAM_SIZE = 2 means to sample bigrams.
     public static final int N_GRAM_SIZE = 2;
 
+    // TODO: Remove dependence on Suggest, and pass in Dictionary as a parameter to an appropriate
+    // method.
     private Suggest mSuggest;
+    @UsedForTesting
+    private Dictionary mDictionaryForTesting;
     private boolean mIsStopping = false;
 
     /* package for test */ int mNumWordsBetweenNGrams;
@@ -83,7 +88,15 @@ public abstract class MainLogBuffer extends FixedLogBuffer {
         mSuggest = suggest;
     }
 
+    @UsedForTesting
+    /* package for test */ void setDictionaryForTesting(final Dictionary dictionary) {
+        mDictionaryForTesting = dictionary;
+    }
+
     private Dictionary getDictionary() {
+        if (mDictionaryForTesting != null) {
+            return mDictionaryForTesting;
+        }
         if (mSuggest == null || !mSuggest.hasMainDictionary()) return null;
         return mSuggest.getMainDictionary();
     }