Add new ContactMatadataProvider.

Bug:20537162
Change-Id: I499efb7c6238957a65c518a1c938533d2fb8ac54
diff --git a/AndroidManifest.xml b/AndroidManifest.xml
index 347abb3..51f4e30 100644
--- a/AndroidManifest.xml
+++ b/AndroidManifest.xml
@@ -71,6 +71,14 @@
             android:permission="com.android.voicemail.permission.ADD_VOICEMAIL">
         </provider>
 
+        <provider android:name="ContactMetadataProvider"
+                  android:authorities="com.android.contacts.metadata"
+                  android:multiprocess="false"
+                  android:exported="true"
+                  android:readPermission="android.permission.READ_CONTACT_METADATA"
+                  android:writePermission="android.permission.WRITE_CONTACT_METADATA">
+        </provider>
+
         <!-- Handles database upgrades after OTAs, then disables itself -->
         <receiver android:name="ContactsUpgradeReceiver">
             <!-- This broadcast is sent after the core system has finished
diff --git a/src/com/android/providers/contacts/ContactMetadataProvider.java b/src/com/android/providers/contacts/ContactMetadataProvider.java
new file mode 100644
index 0000000..31dca7d
--- /dev/null
+++ b/src/com/android/providers/contacts/ContactMetadataProvider.java
@@ -0,0 +1,223 @@
+/*
+ * Copyright (C) 2015 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License
+ */
+package com.android.providers.contacts;
+
+
+import android.content.ContentProvider;
+import android.content.ContentUris;
+import android.content.ContentValues;
+import android.content.Context;
+import android.content.UriMatcher;
+import android.database.Cursor;
+import android.database.DatabaseUtils;
+import android.database.sqlite.SQLiteDatabase;
+import android.database.sqlite.SQLiteQueryBuilder;
+import android.net.Uri;
+import android.os.Binder;
+import android.provider.ContactsContract.MetadataSync;
+import android.text.TextUtils;
+import android.util.Log;
+import com.android.common.content.ProjectionMap;
+import com.android.providers.contacts.ContactsDatabaseHelper.MetadataSyncColumns;
+import com.android.providers.contacts.ContactsDatabaseHelper.Views;
+import com.android.providers.contacts.util.SelectionBuilder;
+import com.android.providers.contacts.util.UserUtils;
+import com.google.common.annotations.VisibleForTesting;
+
+import java.util.Arrays;
+import java.util.Map;
+
+import static com.android.providers.contacts.ContactsProvider2.getLimit;
+import static com.android.providers.contacts.ContactsProvider2.getQueryParameter;
+import static com.android.providers.contacts.util.DbQueryUtils.getEqualityClause;
+
+/**
+ * Simple content provider to handle directing contact metadata specific calls.
+ */
+public class ContactMetadataProvider extends ContentProvider {
+    private static final String TAG = "ContactMetadata";
+    private static final boolean VERBOSE_LOGGING = Log.isLoggable(TAG, Log.VERBOSE);
+    private static final int METADATA_SYNC = 1;
+    private static final int METADATA_SYNC_ID = 2;
+
+    private static final UriMatcher sURIMatcher = new UriMatcher(UriMatcher.NO_MATCH);
+
+    static {
+        sURIMatcher.addURI(MetadataSync.METADATA_AUTHORITY, "metadata_sync", METADATA_SYNC);
+        sURIMatcher.addURI(MetadataSync.METADATA_AUTHORITY, "metadata_sync/#", METADATA_SYNC_ID);
+    }
+
+    private static final Map<String, String> sMetadataProjectionMap = ProjectionMap.builder()
+            .add(MetadataSync._ID)
+            .add(MetadataSync.RAW_CONTACT_BACKUP_ID)
+            .add(MetadataSync.ACCOUNT_TYPE)
+            .add(MetadataSync.ACCOUNT_NAME)
+            .add(MetadataSync.DATA_SET)
+            .add(MetadataSync.DATA)
+            .add(MetadataSync.DELETED)
+            .build();
+
+    private ContactsDatabaseHelper mDbHelper;
+
+    @Override
+    public boolean onCreate() {
+        final Context context = getContext();
+        mDbHelper = getDatabaseHelper(context);
+        return true;
+    }
+
+    protected ContactsDatabaseHelper getDatabaseHelper(final Context context) {
+        return ContactsDatabaseHelper.getInstance(context);
+    }
+
+    @VisibleForTesting
+    protected void setDatabaseHelper(final ContactsDatabaseHelper helper) {
+        mDbHelper = helper;
+    }
+
+    @Override
+    public Cursor query(Uri uri, String[] projection, String selection, String[] selectionArgs,
+            String sortOrder) {
+
+        if (VERBOSE_LOGGING) {
+            Log.v(TAG, "query: uri=" + uri + "  projection=" + Arrays.toString(projection) +
+                    "  selection=[" + selection + "]  args=" + Arrays.toString(selectionArgs) +
+                    "  order=[" + sortOrder + "] CPID=" + Binder.getCallingPid() +
+                    " User=" + UserUtils.getCurrentUserHandle(getContext()));
+        }
+
+        final SQLiteQueryBuilder qb = new SQLiteQueryBuilder();
+        String limit = getLimit(uri);
+        qb.setTables(Views.METADATA_SYNC);
+        qb.setProjectionMap(sMetadataProjectionMap);
+        qb.setStrict(true);
+
+        final SelectionBuilder selectionBuilder = new SelectionBuilder(selection);
+
+        final int match = sURIMatcher.match(uri);
+        switch (match) {
+            case METADATA_SYNC:
+                break;
+
+            case METADATA_SYNC_ID: {
+                selectionBuilder.addClause(getEqualityClause(MetadataSync._ID,
+                        ContentUris.parseId(uri)));
+                break;
+            }
+            default:
+                throw new IllegalArgumentException("Unknown URL " + uri);
+        }
+
+        final SQLiteDatabase db = mDbHelper.getReadableDatabase();
+        return qb.query(db, projection, selectionBuilder.build(), selectionArgs, null,
+                null, sortOrder, limit);
+    }
+
+    @Override
+    public String getType(Uri uri) {
+        int match = sURIMatcher.match(uri);
+        switch (match) {
+            case METADATA_SYNC:
+                return MetadataSync.CONTENT_TYPE;
+            case METADATA_SYNC_ID:
+                return MetadataSync.CONTENT_ITEM_TYPE;
+            default:
+                throw new IllegalArgumentException("Unknown URI: " + uri);
+        }
+    }
+
+    @Override
+    public Uri insert(Uri uri, ContentValues values) {
+
+        if (sURIMatcher.match(uri) != METADATA_SYNC) {
+            throw new IllegalArgumentException(mDbHelper.exceptionMessage(
+                    "Calling contact metadata insert on an unknown/invalid URI" , uri));
+        }
+
+        // Don't insert deleted metadata.
+        Integer deleted = values.getAsInteger(MetadataSync.DELETED);
+        if (deleted != null && deleted != 0) {
+            // Cannot insert deleted metadata
+            throw new IllegalArgumentException(mDbHelper.exceptionMessage(
+                    "Cannot insert deleted metadata:" + values.toString(), uri));
+        }
+
+        // Insert the new entry.
+        // Populate the relevant values before inserting the new entry into the database.
+        final Long accountId = replaceAccountInfoByAccountId(uri, values);
+        final String rawContactBackupId = values.getAsString(MetadataSync.RAW_CONTACT_BACKUP_ID);
+        final String data = values.getAsString(MetadataSync.DATA);
+        deleted = 0; //Only insert non-deleted metadata
+
+        if (accountId == null || rawContactBackupId == null) {
+            throw new IllegalArgumentException(mDbHelper.exceptionMessage(
+                    "Invalid identifier is found: accountId=" + accountId + "; " +
+                            "rawContactBackupId=" + rawContactBackupId, uri));
+        }
+
+        Long metadataSyncId = mDbHelper.replaceMetadataSync(rawContactBackupId, accountId, data,
+                deleted);
+        if (metadataSyncId < 0) {
+            throw new IllegalArgumentException(mDbHelper.exceptionMessage(
+                    "Metadata insertion failed. Values= " + values.toString(), uri));
+        }
+
+        return ContentUris.withAppendedId(uri, metadataSyncId);
+    }
+
+    @Override
+    public int delete(Uri uri, String selection, String[] selectionArgs) {
+        return 0;
+    }
+
+    @Override
+    public int update(Uri uri, ContentValues values, String selection, String[] selectionArgs) {
+        return 0;
+    }
+
+    /**
+     *  Replace account_type, account_name and data_set with account_id. If a valid account_id
+     *  cannot be found for this combination, return null.
+     */
+    private Long replaceAccountInfoByAccountId(Uri uri, ContentValues values) {
+        String accountName = values.getAsString(MetadataSync.ACCOUNT_NAME);
+        String accountType = values.getAsString(MetadataSync.ACCOUNT_TYPE);
+        String dataSet = values.getAsString(MetadataSync.DATA_SET);
+        final boolean partialUri = TextUtils.isEmpty(accountName) ^ TextUtils.isEmpty(accountType);
+        if (partialUri) {
+            // Throw when either account is incomplete.
+            throw new IllegalArgumentException(mDbHelper.exceptionMessage(
+                    "Must specify both or neither of ACCOUNT_NAME and ACCOUNT_TYPE", uri));
+        }
+
+        final AccountWithDataSet account = AccountWithDataSet.get(
+                accountName, accountType, dataSet);
+
+        final Long id = mDbHelper.getAccountIdOrNull(account);
+        if (id == null) {
+            return null;
+        }
+
+        values.put(MetadataSyncColumns.ACCOUNT_ID, id);
+        // Only remove the account information once the account ID is extracted (since these
+        // fields are actually used by resolveAccountWithDataSet to extract the relevant ID).
+        values.remove(MetadataSync.ACCOUNT_NAME);
+        values.remove(MetadataSync.ACCOUNT_TYPE);
+        values.remove(MetadataSync.DATA_SET);
+
+        return id;
+    }
+}
diff --git a/src/com/android/providers/contacts/ContactsDatabaseHelper.java b/src/com/android/providers/contacts/ContactsDatabaseHelper.java
index 2e37e13..59b899e 100644
--- a/src/com/android/providers/contacts/ContactsDatabaseHelper.java
+++ b/src/com/android/providers/contacts/ContactsDatabaseHelper.java
@@ -993,6 +993,7 @@
     private SQLiteStatement mStatusUpdateDelete;
     private SQLiteStatement mResetNameVerifiedForOtherRawContacts;
     private SQLiteStatement mContactInDefaultDirectoryQuery;
+    private SQLiteStatement mMetadataSyncReplace;
 
     private StringBuilder mSb = new StringBuilder();
 
@@ -6034,4 +6035,22 @@
                 " WHERE " + SearchIndexColumns.CONTACT_ID + "=CAST(? AS int)",
                 new String[] {String.valueOf(contactId)});
     }
+
+    public long replaceMetadataSync(String backupId, Long accountId, String data, Integer deleted) {
+        if (mMetadataSyncReplace == null) {
+            mMetadataSyncReplace = getWritableDatabase().compileStatement(
+                    "INSERT OR REPLACE INTO " + Tables.METADATA_SYNC + "("
+                            + MetadataSync.RAW_CONTACT_BACKUP_ID + ", "
+                            + MetadataSyncColumns.ACCOUNT_ID + ", "
+                            + MetadataSync.DATA + ","
+                            + MetadataSync.DELETED + ")" +
+                            " VALUES (?,?,?,?)");
+        }
+        mMetadataSyncReplace.bindString(1, backupId);
+        mMetadataSyncReplace.bindLong(2, accountId);
+        data = (data == null) ? "" : data;
+        mMetadataSyncReplace.bindString(3, data);
+        mMetadataSyncReplace.bindLong(4, deleted);
+        return mMetadataSyncReplace.executeInsert();
+    }
 }
diff --git a/src/com/android/providers/contacts/ContactsProvider2.java b/src/com/android/providers/contacts/ContactsProvider2.java
index 5cbca9e..37633fe 100644
--- a/src/com/android/providers/contacts/ContactsProvider2.java
+++ b/src/com/android/providers/contacts/ContactsProvider2.java
@@ -1973,7 +1973,7 @@
     }
 
     @Override
-    protected ContactsDatabaseHelper getDatabaseHelper(final Context context) {
+    public ContactsDatabaseHelper getDatabaseHelper(final Context context) {
         return ContactsDatabaseHelper.getInstance(context);
     }
 
@@ -8034,7 +8034,7 @@
      * @return A string containing a non-negative integer, or <code>null</code> if
      *         the parameter is not set, or is set to an invalid value.
      */
-    private String getLimit(Uri uri) {
+     static String getLimit(Uri uri) {
         String limitParam = getQueryParameter(uri, ContactsContract.LIMIT_PARAM_KEY);
         if (limitParam == null) {
             return null;
diff --git a/tests/src/com/android/providers/contacts/BaseContactsProvider2Test.java b/tests/src/com/android/providers/contacts/BaseContactsProvider2Test.java
index 3778380..75c9f4b 100644
--- a/tests/src/com/android/providers/contacts/BaseContactsProvider2Test.java
+++ b/tests/src/com/android/providers/contacts/BaseContactsProvider2Test.java
@@ -55,7 +55,7 @@
 import android.test.MoreAsserts;
 import android.test.mock.MockContentResolver;
 import android.util.Log;
-
+import com.android.providers.contacts.ContactsDatabaseHelper.AccountsColumns;
 import com.android.providers.contacts.ContactsDatabaseHelper.Tables;
 import com.android.providers.contacts.testutil.CommonDatabaseUtils;
 import com.android.providers.contacts.testutil.DataUtil;
@@ -133,6 +133,8 @@
         mActor.addPermissions(
                 "android.permission.READ_CONTACTS",
                 "android.permission.WRITE_CONTACTS",
+                "android.permission.READ_CONTACT_METADATA",
+                "android.permission.WRITE_CONTACT_METADATA",
                 "android.permission.READ_SOCIAL_STREAM",
                 "android.permission.WRITE_SOCIAL_STREAM",
                 "android.permission.READ_PROFILE",
@@ -563,6 +565,18 @@
                 " WHERE " + BaseColumns._ID + "=" + contactId);
     }
 
+    protected long createAccount(String accountName, String accountType, String dataSet) {
+        // There's no api for this, so we just tweak the DB directly.
+        SQLiteDatabase db = ((ContactsProvider2) getProvider()).getDatabaseHelper()
+                .getWritableDatabase();
+
+        ContentValues values = new ContentValues();
+        values.put(AccountsColumns.ACCOUNT_NAME, accountName);
+        values.put(AccountsColumns.ACCOUNT_TYPE, accountType);
+        values.put(AccountsColumns.DATA_SET, dataSet);
+        return db.insert(Tables.ACCOUNTS, null, values);
+    }
+
     protected Cursor queryRawContact(long rawContactId) {
         return mResolver.query(ContentUris.withAppendedId(RawContacts.CONTENT_URI, rawContactId),
                 null, null, null, null);
diff --git a/tests/src/com/android/providers/contacts/ContactMetadataProviderTest.java b/tests/src/com/android/providers/contacts/ContactMetadataProviderTest.java
new file mode 100644
index 0000000..7f45ac8
--- /dev/null
+++ b/tests/src/com/android/providers/contacts/ContactMetadataProviderTest.java
@@ -0,0 +1,143 @@
+/*
+ * Copyright (C) 2015 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.providers.contacts;
+
+import android.content.ContentValues;
+import android.database.Cursor;
+import android.net.Uri;
+import android.provider.ContactsContract;
+import android.provider.ContactsContract.Data;
+import android.provider.ContactsContract.MetadataSync;
+import android.provider.ContactsContract.RawContacts;
+import android.test.suitebuilder.annotation.MediumTest;
+import com.android.providers.contacts.ContactsDatabaseHelper.MetadataSyncColumns;
+import com.android.providers.contacts.testutil.RawContactUtil;
+import com.android.providers.contacts.testutil.TestUtil;
+
+/**
+ * Unit tests for {@link com.android.providers.contacts.ContactMetadataProvider}.
+ * <p/>
+ * Run the test like this:
+ * <code>
+ * adb shell am instrument -e class com.android.providers.contacts.ContactMetadataProviderTest -w \
+ * com.android.providers.contacts.tests/android.test.InstrumentationTestRunner
+ * </code>
+ */
+@MediumTest
+public class ContactMetadataProviderTest extends BaseContactsProvider2Test {
+    private static String TEST_ACCOUNT_TYPE = "test_account_type";
+    private static String TEST_ACCOUNT_NAME = "test_account_name";
+    private static String TEST_DATA_SET = "test_data_set";
+    private static String TEST_BACKUP_ID = "1001";
+    private static String TEST_DATA = "test_data";
+    private static String SELECTION_BY_TEST_ACCOUNT = MetadataSync.ACCOUNT_NAME + "='" +
+            TEST_ACCOUNT_NAME + "' AND " + MetadataSync.ACCOUNT_TYPE + "='" + TEST_ACCOUNT_TYPE +
+            "' AND " + MetadataSync.DATA_SET + "='" + TEST_DATA_SET + "'";
+
+    private ContactMetadataProvider mContactMetadataProvider;
+    private AccountWithDataSet mTestAccount;
+    private ContentValues defaultValues;
+
+    @Override
+    protected void setUp() throws Exception {
+        super.setUp();
+        mContactMetadataProvider = (ContactMetadataProvider) addProvider(
+               ContactMetadataProvider.class, MetadataSync.METADATA_AUTHORITY);
+        // Reset the dbHelper to be the one ContactsProvider2 is using. Before this, two providers
+        // are using different dbHelpers.
+        mContactMetadataProvider.setDatabaseHelper(((SynchronousContactsProvider2)
+                mActor.provider).getDatabaseHelper(getContext()));
+        setupData();
+    }
+
+    public void testInsertWithInvalidUri() {
+        try {
+            mResolver.insert(Uri.withAppendedPath(MetadataSync.METADATA_AUTHORITY_URI,
+                    "metadata"), getDefaultValues());
+            fail("the insert was expected to fail, but it succeeded");
+        } catch (IllegalArgumentException e) {
+            // this was expected
+        }
+    }
+
+    public void testGetMetadataByAccount() {
+        Cursor c = mResolver.query(MetadataSync.CONTENT_URI, null, SELECTION_BY_TEST_ACCOUNT,
+                null, null);
+        assertEquals(1, c.getCount());
+
+        ContentValues expectedValues = defaultValues;
+        expectedValues.remove(MetadataSyncColumns.ACCOUNT_ID);
+        c.moveToFirst();
+        assertCursorValues(c, expectedValues);
+        c.close();
+    }
+
+    public void testReplaceMetadataForSameAccountIdAndBackupId() {
+        //Insert a new metadata with same account and backupId as defaultValues, but different data
+        //field.
+        ContentValues  newValues =  new ContentValues();
+        newValues.put(MetadataSync.ACCOUNT_NAME, TEST_ACCOUNT_NAME);
+        newValues.put(MetadataSync.ACCOUNT_TYPE, TEST_ACCOUNT_TYPE);
+        newValues.put(MetadataSync.DATA_SET, TEST_DATA_SET);
+        newValues.put(MetadataSync.RAW_CONTACT_BACKUP_ID, TEST_BACKUP_ID);
+        newValues.put(MetadataSync.DATA, "new data");
+        newValues.put(MetadataSync.DELETED, 0);
+        mResolver.insert(MetadataSync.CONTENT_URI, newValues);
+
+        // Total two metadata entries
+        Cursor c = mResolver.query(MetadataSync.CONTENT_URI, null, null,
+                null, null);
+        assertEquals(2, c.getCount());
+
+        // Only one metadata entry for TEST_ACCOUNT
+        c = mResolver.query(MetadataSync.CONTENT_URI, null, SELECTION_BY_TEST_ACCOUNT, null, null);
+        assertEquals(1, c.getCount());
+        newValues.remove(MetadataSyncColumns.ACCOUNT_ID);
+        c.moveToFirst();
+        assertCursorValues(c, newValues);
+        c.close();
+    }
+
+    private void setupData() {
+        mTestAccount = new AccountWithDataSet(TEST_ACCOUNT_NAME, TEST_ACCOUNT_TYPE, TEST_DATA_SET);
+        long rawContactId1 = RawContactUtil.createRawContactWithAccountDataSet(
+                mResolver, mTestAccount);
+        createAccount(TEST_ACCOUNT_NAME, TEST_ACCOUNT_TYPE, TEST_DATA_SET);
+        mResolver.insert(MetadataSync.CONTENT_URI, getDefaultValues());
+
+        // Insert another entry for another account
+        createAccount("John", "account2", null);
+        ContentValues values = new ContentValues();
+        values.put(MetadataSync.ACCOUNT_NAME, "John");
+        values.put(MetadataSync.ACCOUNT_TYPE, "account2");
+        values.put(MetadataSync.RAW_CONTACT_BACKUP_ID, "1");
+        values.put(MetadataSync.DATA, TEST_DATA);
+        values.put(MetadataSync.DELETED, 0);
+        mResolver.insert(MetadataSync.CONTENT_URI, values);
+    }
+
+    private ContentValues getDefaultValues() {
+        defaultValues = new ContentValues();
+        defaultValues.put(MetadataSync.ACCOUNT_NAME, TEST_ACCOUNT_NAME);
+        defaultValues.put(MetadataSync.ACCOUNT_TYPE, TEST_ACCOUNT_TYPE);
+        defaultValues.put(MetadataSync.DATA_SET, TEST_DATA_SET);
+        defaultValues.put(MetadataSync.RAW_CONTACT_BACKUP_ID, TEST_BACKUP_ID);
+        defaultValues.put(MetadataSync.DATA, TEST_DATA);
+        defaultValues.put(MetadataSync.DELETED, 0);
+        return defaultValues;
+    }
+}
diff --git a/tests/src/com/android/providers/contacts/StandaloneContactsProvider2.java b/tests/src/com/android/providers/contacts/StandaloneContactsProvider2.java
index a55072c..8dd09bf 100644
--- a/tests/src/com/android/providers/contacts/StandaloneContactsProvider2.java
+++ b/tests/src/com/android/providers/contacts/StandaloneContactsProvider2.java
@@ -29,7 +29,7 @@
     }
 
     @Override
-    protected ContactsDatabaseHelper getDatabaseHelper(final Context context) {
+    public ContactsDatabaseHelper getDatabaseHelper(final Context context) {
         if (mDbHelper == null) {
             mDbHelper = ContactsDatabaseHelper.getNewInstanceForTest(context);
         }
diff --git a/tests/src/com/android/providers/contacts/SynchronousContactsProvider2.java b/tests/src/com/android/providers/contacts/SynchronousContactsProvider2.java
index 1d127c7..e24cd00 100644
--- a/tests/src/com/android/providers/contacts/SynchronousContactsProvider2.java
+++ b/tests/src/com/android/providers/contacts/SynchronousContactsProvider2.java
@@ -41,7 +41,7 @@
     private boolean mIsVoiceCapable = true;
 
     @Override
-    protected ContactsDatabaseHelper getDatabaseHelper(final Context context) {
+    public ContactsDatabaseHelper getDatabaseHelper(final Context context) {
         if (sDbHelper == null) {
             sDbHelper = ContactsDatabaseHelper.getNewInstanceForTest(context);
         }
diff --git a/tests/src/com/android/providers/contacts/testutil/RawContactUtil.java b/tests/src/com/android/providers/contacts/testutil/RawContactUtil.java
index e9cd3b5..5b38be9 100644
--- a/tests/src/com/android/providers/contacts/testutil/RawContactUtil.java
+++ b/tests/src/com/android/providers/contacts/testutil/RawContactUtil.java
@@ -24,6 +24,7 @@
 import android.net.Uri;
 import android.provider.ContactsContract;
 import android.test.mock.MockContentResolver;
+import com.android.providers.contacts.AccountWithDataSet;
 
 import java.util.List;
 
@@ -90,7 +91,18 @@
             String... extras) {
         ContentValues values = new ContentValues();
         CommonDatabaseUtils.extrasVarArgsToValues(values, extras);
-        final Uri uri = TestUtil.maybeAddAccountQueryParameters(ContactsContract.RawContacts.CONTENT_URI, account);
+        final Uri uri = TestUtil.maybeAddAccountQueryParameters(
+                ContactsContract.RawContacts.CONTENT_URI, account);
+        Uri contactUri = resolver.insert(uri, values);
+        return ContentUris.parseId(contactUri);
+    }
+
+    public static long createRawContactWithAccountDataSet(ContentResolver resolver,
+            AccountWithDataSet accountWithDataSet, String... extras) {
+        ContentValues values = new ContentValues();
+        CommonDatabaseUtils.extrasVarArgsToValues(values, extras);
+        final Uri uri = TestUtil.maybeAddAccountWithDataSetQueryParameters(
+                ContactsContract.RawContacts.CONTENT_URI, accountWithDataSet);
         Uri contactUri = resolver.insert(uri, values);
         return ContentUris.parseId(contactUri);
     }
diff --git a/tests/src/com/android/providers/contacts/testutil/TestUtil.java b/tests/src/com/android/providers/contacts/testutil/TestUtil.java
index 2020f6d..05ff61d 100644
--- a/tests/src/com/android/providers/contacts/testutil/TestUtil.java
+++ b/tests/src/com/android/providers/contacts/testutil/TestUtil.java
@@ -18,8 +18,9 @@
 
 import android.accounts.Account;
 import android.net.Uri;
-import android.provider.ContactsContract;
+import android.provider.ContactsContract.RawContacts;
 import android.util.Log;
+import com.android.providers.contacts.AccountWithDataSet;
 
 /**
  * Common methods used for testing.
@@ -46,8 +47,20 @@
             return uri;
         }
         return uri.buildUpon()
-                .appendQueryParameter(ContactsContract.RawContacts.ACCOUNT_NAME, account.name)
-                .appendQueryParameter(ContactsContract.RawContacts.ACCOUNT_TYPE, account.type)
+                .appendQueryParameter(RawContacts.ACCOUNT_NAME, account.name)
+                .appendQueryParameter(RawContacts.ACCOUNT_TYPE, account.type)
+                .build();
+    }
+
+    public static Uri maybeAddAccountWithDataSetQueryParameters(Uri uri,
+            AccountWithDataSet account) {
+        if (account == null) {
+            return uri;
+        }
+        return uri.buildUpon()
+                .appendQueryParameter(RawContacts.ACCOUNT_NAME, account.getAccountName())
+                .appendQueryParameter(RawContacts.ACCOUNT_TYPE, account.getAccountType())
+                .appendQueryParameter(RawContacts.DATA_SET, account.getDataSet())
                 .build();
     }
 }