diff --git a/tests/src/com/android/inputmethod/latin/SuggestHelper.java b/tests/src/com/android/inputmethod/latin/SuggestHelper.java
index cccd1a4..2e36245 100644
--- a/tests/src/com/android/inputmethod/latin/SuggestHelper.java
+++ b/tests/src/com/android/inputmethod/latin/SuggestHelper.java
@@ -21,7 +21,7 @@
 
 import com.android.inputmethod.keyboard.KeyDetector;
 import com.android.inputmethod.keyboard.Keyboard;
-import com.android.inputmethod.keyboard.KeyboardId;
+import com.android.inputmethod.keyboard.KeyboardSet;
 
 import java.io.File;
 import java.util.Locale;
@@ -32,24 +32,20 @@
     protected final Keyboard mKeyboard;
     private final KeyDetector mKeyDetector;
 
-    public static final int ALPHABET_KEYBOARD = com.android.inputmethod.latin.R.xml.kbd_qwerty;
-
-    public SuggestHelper(Context context, int dictionaryId, KeyboardId keyboardId) {
+    public SuggestHelper(Context context, int dictionaryId, KeyboardSet keyboardSet) {
         // Use null as the locale for Suggest so as to force it to use the internal dictionary
         // (and not try to find a dictionary provider for a specified locale)
         mSuggest = new Suggest(context, dictionaryId, null);
-        mKeyboard = new Keyboard.Builder<Keyboard.Params>(context, new Keyboard.Params())
-                .load(ALPHABET_KEYBOARD, keyboardId).build();
+        mKeyboard = keyboardSet.getMainKeyboard();
         mKeyDetector = new KeyDetector(0);
         init();
     }
 
     protected SuggestHelper(final Context context, final File dictionaryPath,
-            final long startOffset, final long length, final KeyboardId keyboardId,
+            final long startOffset, final long length, final KeyboardSet keyboardSet,
             final Locale locale) {
         mSuggest = new Suggest(context, dictionaryPath, startOffset, length, null, locale);
-        mKeyboard = new Keyboard.Builder<Keyboard.Params>(context, new Keyboard.Params())
-                .load(ALPHABET_KEYBOARD, keyboardId).build();
+        mKeyboard = keyboardSet.getMainKeyboard();
         mKeyDetector = new KeyDetector(0);
         init();
     }
diff --git a/tests/src/com/android/inputmethod/latin/SuggestTests.java b/tests/src/com/android/inputmethod/latin/SuggestTests.java
index 5b1c9fa..e12ae58 100644
--- a/tests/src/com/android/inputmethod/latin/SuggestTests.java
+++ b/tests/src/com/android/inputmethod/latin/SuggestTests.java
@@ -33,7 +33,7 @@
         final Locale locale = Locale.US;
         mHelper = new SuggestHelper(
                 getContext(), mTestPackageFile, dict.getStartOffset(), dict.getLength(),
-                createKeyboardId(locale, Configuration.ORIENTATION_PORTRAIT), locale);
+                createKeyboardSet(locale, Configuration.ORIENTATION_PORTRAIT), locale);
         mHelper.setCorrectionMode(Suggest.CORRECTION_FULL_BIGRAM);
     }
 
diff --git a/tests/src/com/android/inputmethod/latin/SuggestTestsBase.java b/tests/src/com/android/inputmethod/latin/SuggestTestsBase.java
index 9dd61d7..73e34ba 100644
--- a/tests/src/com/android/inputmethod/latin/SuggestTestsBase.java
+++ b/tests/src/com/android/inputmethod/latin/SuggestTestsBase.java
@@ -22,8 +22,9 @@
 import android.text.InputType;
 import android.text.TextUtils;
 import android.util.DisplayMetrics;
+import android.view.inputmethod.EditorInfo;
 
-import com.android.inputmethod.keyboard.KeyboardId;
+import com.android.inputmethod.keyboard.KeyboardSet;
 
 import java.io.File;
 import java.io.InputStream;
@@ -38,7 +39,12 @@
         mTestPackageFile = new File(getTestContext().getApplicationInfo().sourceDir);
     }
 
-    protected KeyboardId createKeyboardId(Locale locale, int orientation) {
+    protected KeyboardSet createKeyboardSet(Locale locale, int orientation) {
+        return createKeyboardSet(locale, orientation, false);
+    }
+
+    protected KeyboardSet createKeyboardSet(Locale locale, int orientation,
+            boolean touchPositionCorrectionEnabled) {
         final DisplayMetrics dm = getContext().getResources().getDisplayMetrics();
         final int width;
         if (orientation == Configuration.ORIENTATION_LANDSCAPE) {
@@ -50,8 +56,12 @@
                     + "orientation=" + orientation);
             return null;
         }
-        return new KeyboardId(KeyboardId.ELEMENT_ALPHABET, locale, orientation, width,
-                KeyboardId.MODE_TEXT, InputType.TYPE_CLASS_TEXT, 0, false, false, false, false);
+        final EditorInfo editorInfo = new EditorInfo();
+        editorInfo.inputType = InputType.TYPE_CLASS_TEXT;
+        final KeyboardSet.Builder builder = new KeyboardSet.Builder(getContext(), editorInfo);
+        builder.setScreenGeometry(orientation, width);
+        builder.setSubtype(locale, true, touchPositionCorrectionEnabled);
+        return builder.build();
     }
 
     protected InputStream openTestRawResource(int resIdInTest) {
diff --git a/tests/src/com/android/inputmethod/latin/UserBigramSuggestHelper.java b/tests/src/com/android/inputmethod/latin/UserBigramSuggestHelper.java
index 863c2b2..74fadf7 100644
--- a/tests/src/com/android/inputmethod/latin/UserBigramSuggestHelper.java
+++ b/tests/src/com/android/inputmethod/latin/UserBigramSuggestHelper.java
@@ -16,11 +16,11 @@
 
 package com.android.inputmethod.latin;
 
-import com.android.inputmethod.keyboard.KeyboardId;
-
 import android.content.Context;
 import android.text.TextUtils;
 
+import com.android.inputmethod.keyboard.KeyboardSet;
+
 import java.io.File;
 import java.util.Locale;
 import java.util.StringTokenizer;
@@ -31,8 +31,8 @@
 
     public UserBigramSuggestHelper(final Context context, final File dictionaryPath,
             final long startOffset, final long length, final int userBigramMax,
-            final int userBigramDelete, final KeyboardId keyboardId, final Locale locale) {
-        super(context, dictionaryPath, startOffset, length, keyboardId, locale);
+            final int userBigramDelete, final KeyboardSet keyboardSet, final Locale locale) {
+        super(context, dictionaryPath, startOffset, length, keyboardSet, locale);
         mContext = context;
         mUserBigram = new UserBigramDictionary(context, null, locale.toString(),
                 Suggest.DIC_USER);
diff --git a/tests/src/com/android/inputmethod/latin/UserBigramSuggestTests.java b/tests/src/com/android/inputmethod/latin/UserBigramSuggestTests.java
index b7ba9c7..2b88a7c 100644
--- a/tests/src/com/android/inputmethod/latin/UserBigramSuggestTests.java
+++ b/tests/src/com/android/inputmethod/latin/UserBigramSuggestTests.java
@@ -37,7 +37,7 @@
         mHelper = new UserBigramSuggestHelper(
                 getContext(), mTestPackageFile, dict.getStartOffset(), dict.getLength(),
                 MAX_DATA, DELETE_DATA,
-                createKeyboardId(locale, Configuration.ORIENTATION_PORTRAIT), locale);
+                createKeyboardSet(locale, Configuration.ORIENTATION_PORTRAIT), locale);
     }
 
     /************************** Tests ************************/