Replace cookie setting toggles with four state radio buttons

Begin cookie settings redesign by replacing the two toggles with a
four-state preference, hidden behind a flag. Adjust tests to use the new
four-state preference.
https://screenshot.googleplex.com/r22q4q6hH7O.png

Bug: 1060118
Change-Id: I5f660b473112e8ab688b8ac5604f8bcceb7f35a0
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2095124
Commit-Queue: Ehimare Okoyomon <eokoyomon@chromium.org>
Reviewed-by: Finnur Thorarinsson <finnur@chromium.org>
Reviewed-by: Christian Dullweber <dullweber@chromium.org>
Reviewed-by: Natalie Chouinard <chouinard@chromium.org>
Reviewed-by: Theresa  <twellington@chromium.org>
Cr-Commit-Position: refs/heads/master@{#749638}

chrome/android/chrome_java_sources.gni

@@ -1484,6 +1484,7 @@ chrome_java_sources = [
   "java/src/org/chromium/chrome/browser/site_settings/ContentSettingsResources.java",
   "java/src/org/chromium/chrome/browser/site_settings/CookieControlsBridge.java",
   "java/src/org/chromium/chrome/browser/site_settings/CookieControlsServiceBridge.java",
+  "java/src/org/chromium/chrome/browser/site_settings/FourStateCookieSettingsPreference.java",
   "java/src/org/chromium/chrome/browser/site_settings/LocalStorageInfo.java",
   "java/src/org/chromium/chrome/browser/site_settings/LocationCategory.java",
   "java/src/org/chromium/chrome/browser/site_settings/ManageSpaceActivity.java",

chrome/android/java/res/layout/four_state_cookie_settings_preference.xml

+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright 2020 The Chromium Authors. All rights reserved.
+     Use of this source code is governed by a BSD-style license that can be
+     found in the LICENSE file. -->
+
+<!-- Layout used by the FourStateCookieSettingsPreference. -->
+
+<FrameLayout
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    xmlns:app="http://schemas.android.com/apk/res-auto"
+    android:layout_width="match_parent"
+    android:layout_height="wrap_content"
+    android:focusable="false"
+    android:orientation="vertical">
+
+    <org.chromium.components.browser_ui.widget.RadioButtonWithDescriptionLayout
+        android:id="@+id/radio_button_layout"
+        android:layout_width="match_parent"
+        android:layout_height="match_parent">
+
+        <org.chromium.components.browser_ui.widget.RadioButtonWithDescription
+            android:id="@+id/allow"
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            app:primaryText="@string/website_settings_category_cookie_allow_title"
+            app:descriptionText="@string/website_settings_category_cookie_allow_addition" />
+
+        <org.chromium.components.browser_ui.widget.RadioButtonWithDescription
+            android:id="@+id/block_third_party_incognito"
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            app:primaryText="@string/website_settings_category_cookie_block_third_party_incognito_title"
+            app:descriptionText="@string/website_settings_category_cookie_block_third_party_addition" />
+
+        <org.chromium.components.browser_ui.widget.RadioButtonWithDescription
+            android:id="@+id/block_third_party"
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            app:primaryText="@string/website_settings_category_cookie_block_third_party_title"
+            app:descriptionText="@string/website_settings_category_cookie_block_third_party_addition" />
+
+        <org.chromium.components.browser_ui.widget.RadioButtonWithDescription
+            android:id="@+id/block"
+            android:layout_width="match_parent"
+            android:layout_height="wrap_content"
+            app:primaryText="@string/website_settings_category_cookie_block_title"
+            app:descriptionText="@string/website_settings_category_cookie_block_addition" />
+
+    </org.chromium.components.browser_ui.widget.RadioButtonWithDescriptionLayout>
+</FrameLayout>

chrome/android/java/res/xml/website_preferences.xml

@@ -19,10 +19,13 @@
         android:key="binary_toggle"
         android:defaultValue="true" />
     <!-- A common Allow/Ask/Block 3-state toggle (radio group). Only shown when
-         the category requires these 3 states. Mutual exclusive with the
-         "binary_toggle" above. -->
+         the category requires these 3 states. Mutually exclusive with the
+         "binary_toggle" above and the "four_state_cookie_toggle" below. -->
     <org.chromium.chrome.browser.site_settings.TriStateSiteSettingsPreference
         android:key="tri_state_toggle" />
+    <!-- A four state toggle for cookie preferences, only shown for the Cookies category. -->
+    <org.chromium.chrome.browser.site_settings.FourStateCookieSettingsPreference
+        android:key="four_state_cookie_toggle" />
 
     <!-- A toggle for blocking third-party cookies, only shown for the Cookies category. -->
     <org.chromium.chrome.browser.settings.ChromeBaseCheckBoxPreference

chrome/android/java/src/org/chromium/chrome/browser/site_settings/FourStateCookieSettingsPreference.java

+// Copyright 2020 The Chromium Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+package org.chromium.chrome.browser.site_settings;
+
+import android.content.Context;
+import android.support.v7.preference.Preference;
+import android.support.v7.preference.PreferenceViewHolder;
+import android.util.AttributeSet;
+import android.widget.RadioGroup;
+
+import org.chromium.chrome.R;
+import org.chromium.components.browser_ui.widget.RadioButtonWithDescription;
+
+/**
+ * A 4-state radio group Preference used for the Cookies subpage of SiteSettings.
+ */
+public class FourStateCookieSettingsPreference
+        extends Preference implements RadioGroup.OnCheckedChangeListener {
+    public enum CookieSettingsState {
+        UNINITIALIZED,
+        ALLOW,
+        BLOCK_THIRD_PARTY_INCOGNITO,
+        BLOCK_THIRD_PARTY,
+        BLOCK
+    }
+
+    private CookieSettingsState mState = CookieSettingsState.UNINITIALIZED;
+    private RadioButtonWithDescription mAllow;
+    private RadioButtonWithDescription mBlockThirdPartyIncognito;
+    private RadioButtonWithDescription mBlockThirdParty;
+    private RadioButtonWithDescription mBlock;
+    private RadioGroup mRadioGroup;
+
+    public FourStateCookieSettingsPreference(Context context, AttributeSet attrs) {
+        super(context, attrs);
+
+        // Sets the layout resource that will be inflated for the view.
+        setLayoutResource(R.layout.four_state_cookie_settings_preference);
+
+        // Make unselectable, otherwise FourStateCookieSettingsPreference is treated as one
+        // selectable Preference, instead of four selectable radio buttons.
+        setSelectable(false);
+    }
+
+    /**
+     * Sets the state of this Preference.
+     * @param state The CookieSettingsState to set the preference to.
+     */
+    public void setState(CookieSettingsState state) {
+        mState = state;
+
+        if (mRadioGroup != null) {
+            setRadioButton();
+        }
+    }
+
+    /**
+     * @return The state that is currently selected.
+     */
+    public CookieSettingsState getState() {
+        return mState;
+    }
+
+    @Override
+    public void onCheckedChanged(RadioGroup group, int checkedId) {
+        if (mAllow.isChecked()) {
+            mState = CookieSettingsState.ALLOW;
+        } else if (mBlockThirdPartyIncognito.isChecked()) {
+            mState = CookieSettingsState.BLOCK_THIRD_PARTY_INCOGNITO;
+        } else if (mBlockThirdParty.isChecked()) {
+            mState = CookieSettingsState.BLOCK_THIRD_PARTY;
+        } else if (mBlock.isChecked()) {
+            mState = CookieSettingsState.BLOCK;
+        }
+
+        callChangeListener(mState);
+    }
+
+    @Override
+    public void onBindViewHolder(PreferenceViewHolder holder) {
+        super.onBindViewHolder(holder);
+
+        mAllow = (RadioButtonWithDescription) holder.findViewById(R.id.allow);
+        mBlockThirdPartyIncognito =
+                (RadioButtonWithDescription) holder.findViewById(R.id.block_third_party_incognito);
+        mBlockThirdParty = (RadioButtonWithDescription) holder.findViewById(R.id.block_third_party);
+        mBlock = (RadioButtonWithDescription) holder.findViewById(R.id.block);
+        mRadioGroup = (RadioGroup) holder.findViewById(R.id.radio_button_layout);
+        mRadioGroup.setOnCheckedChangeListener(this);
+
+        if (mState != CookieSettingsState.UNINITIALIZED) {
+            setRadioButton();
+        }
+    }
+
+    private void setRadioButton() {
+        RadioButtonWithDescription button = null;
+
+        switch (mState) {
+            case ALLOW:
+                button = mAllow;
+                break;
+            case BLOCK_THIRD_PARTY_INCOGNITO:
+                button = mBlockThirdPartyIncognito;
+                break;
+            case BLOCK_THIRD_PARTY:
+                button = mBlockThirdParty;
+                break;
+            case BLOCK:
+                button = mBlock;
+                break;
+            default:
+                return;
+        }
+
+        button.setChecked(true);
+    }
+}

chrome/android/java/src/org/chromium/chrome/browser/site_settings/SingleCategorySettings.java

@@ -52,8 +52,10 @@ import org.chromium.chrome.browser.settings.ManagedPreferenceDelegate;
 import org.chromium.chrome.browser.settings.ManagedPreferencesUtils;
 import org.chromium.chrome.browser.settings.SearchUtils;
 import org.chromium.chrome.browser.settings.SettingsUtils;
+import org.chromium.chrome.browser.site_settings.FourStateCookieSettingsPreference.CookieSettingsState;
 import org.chromium.chrome.browser.site_settings.Website.StoredDataClearedCallback;
 import org.chromium.components.content_settings.ContentSettingsType;
+import org.chromium.components.content_settings.CookieControlsMode;
 import org.chromium.components.embedder_support.util.UrlUtilities;
 import org.chromium.ui.widget.Toast;
 
@@ -116,6 +118,8 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
     private List<WebsitePreference> mWebsites;
     // Whether tri-state ContentSetting is required.
     private boolean mRequiresTriStateSetting;
+    // Whether four-state ContentSetting is required.
+    private boolean mRequiresFourStateSetting;
     // Locally-saved reference to the "notifications_quiet_ui" preference to allow hiding/showing
     // it.
     private ChromeBaseCheckBoxPreference mNotificationsQuietUiPref;
@@ -123,10 +127,11 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
     @Nullable
     private Set<String> mSelectedDomains;
 
-    // Keys for common ContentSetting toggle for categories. These two toggles are mutually
+    // Keys for common ContentSetting toggle for categories. These three toggles are mutually
     // exclusive: a category should only show one of them, at most.
     public static final String BINARY_TOGGLE_KEY = "binary_toggle";
     public static final String TRI_STATE_TOGGLE_KEY = "tri_state_toggle";
+    public static final String FOUR_STATE_COOKIE_TOGGLE_KEY = "four_state_cookie_toggle";
 
     // Keys for category-specific preferences (toggle, link, button etc.), dynamically shown.
     public static final String THIRD_PARTY_COOKIES_TOGGLE_KEY = "third_party_cookies";
@@ -290,6 +295,8 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
         int contentType = mCategory.getContentSettingsType();
         mRequiresTriStateSetting =
                 WebsitePreferenceBridge.requiresTriStateContentSetting(contentType);
+        mRequiresFourStateSetting =
+                WebsitePreferenceBridge.requiresFourStateContentSetting(contentType);
 
         ViewGroup view = (ViewGroup) super.onCreateView(inflater, container, savedInstanceState);
 
@@ -517,6 +524,8 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
             int setting = (int) newValue;
             WebsitePreferenceBridge.setContentSetting(mCategory.getContentSettingsType(), setting);
             getInfoForOrigins();
+        } else if (FOUR_STATE_COOKIE_TOGGLE_KEY.equals(preference.getKey())) {
+            setCookieSettingsPreference((CookieSettingsState) newValue);
         } else if (THIRD_PARTY_COOKIES_TOGGLE_KEY.equals(preference.getKey())) {
             PrefServiceBridge.getInstance().setBoolean(
                     Pref.BLOCK_THIRD_PARTY_COOKIES, ((boolean) newValue));
@@ -529,6 +538,38 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
         return true;
     }
 
+    private void setCookieSettingsPreference(CookieSettingsState state) {
+        boolean allowCookies;
+        @CookieControlsMode
+        int mode;
+
+        switch (state) {
+            case ALLOW:
+                allowCookies = true;
+                mode = CookieControlsMode.OFF;
+                break;
+            case BLOCK_THIRD_PARTY_INCOGNITO:
+                allowCookies = true;
+                mode = CookieControlsMode.INCOGNITO_ONLY;
+                break;
+            case BLOCK_THIRD_PARTY:
+                allowCookies = true;
+                mode = CookieControlsMode.ON;
+                break;
+            case BLOCK:
+                allowCookies = false;
+                mode = CookieControlsMode.ON;
+                break;
+            default:
+                return;
+        }
+
+        WebsitePreferenceBridge.setCategoryEnabled(ContentSettingsType.COOKIES, allowCookies);
+        PrefServiceBridge.getInstance().setInteger(Pref.COOKIE_CONTROLS_MODE, mode);
+        PrefServiceBridge.getInstance().setBoolean(
+                Pref.BLOCK_THIRD_PARTY_COOKIES, mode == CookieControlsMode.ON);
+    }
+
     private String getAddExceptionDialogMessage() {
         int resource = 0;
         if (mCategory.showSites(SiteSettingsCategory.Type.AUTOMATIC_DOWNLOADS)) {
@@ -810,6 +851,11 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
                     (TriStateSiteSettingsPreference) getPreferenceScreen().findPreference(
                             TRI_STATE_TOGGLE_KEY);
             return (triStateToggle.getCheckedSetting() == ContentSettingValues.BLOCK);
+        } else if (mRequiresFourStateSetting) {
+            FourStateCookieSettingsPreference fourStateCookieToggle =
+                    (FourStateCookieSettingsPreference) getPreferenceScreen().findPreference(
+                            FOUR_STATE_COOKIE_TOGGLE_KEY);
+            return fourStateCookieToggle.getState() == CookieSettingsState.BLOCK;
         } else {
             ChromeSwitchPreference binaryToggle =
                     (ChromeSwitchPreference) getPreferenceScreen().findPreference(
@@ -829,6 +875,9 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
                 (ChromeSwitchPreference) screen.findPreference(BINARY_TOGGLE_KEY);
         TriStateSiteSettingsPreference triStateToggle =
                 (TriStateSiteSettingsPreference) screen.findPreference(TRI_STATE_TOGGLE_KEY);
+        FourStateCookieSettingsPreference fourStateCookieToggle =
+                (FourStateCookieSettingsPreference) screen.findPreference(
+                        FOUR_STATE_COOKIE_TOGGLE_KEY);
         Preference thirdPartyCookies = screen.findPreference(THIRD_PARTY_COOKIES_TOGGLE_KEY);
         Preference notificationsVibrate = screen.findPreference(NOTIFICATIONS_VIBRATE_TOGGLE_KEY);
         Preference notificationsQuietUi = screen.findPreference(NOTIFICATIONS_QUIET_UI_TOGGLE_KEY);
@@ -848,11 +897,18 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
         if (hideMainToggles) {
             screen.removePreference(binaryToggle);
             screen.removePreference(triStateToggle);
+            screen.removePreference(fourStateCookieToggle);
         } else if (mRequiresTriStateSetting) {
             screen.removePreference(binaryToggle);
+            screen.removePreference(fourStateCookieToggle);
             configureTriStateToggle(triStateToggle, contentType);
+        } else if (mRequiresFourStateSetting) {
+            screen.removePreference(binaryToggle);
+            screen.removePreference(triStateToggle);
+            configureFourStateCookieToggle(fourStateCookieToggle);
         } else {
             screen.removePreference(triStateToggle);
+            screen.removePreference(fourStateCookieToggle);
             configureBinaryToggle(binaryToggle, contentType);
         }
 
@@ -879,7 +935,7 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
         }
 
         // Configure/hide the third-party cookies toggle, as needed.
-        if (mCategory.showSites(SiteSettingsCategory.Type.COOKIES)) {
+        if (mCategory.showSites(SiteSettingsCategory.Type.COOKIES) && !mRequiresFourStateSetting) {
             thirdPartyCookies.setOnPreferenceChangeListener(this);
             updateThirdPartyCookiesCheckBox();
         } else {
@@ -953,6 +1009,34 @@ public class SingleCategorySettings extends PreferenceFragmentCompat
         }
     }
 
+    private void configureFourStateCookieToggle(
+            FourStateCookieSettingsPreference fourStateCookieToggle) {
+        fourStateCookieToggle.setOnPreferenceChangeListener(this);
+        // These conditions only check the preference combinations that deterministically decide
+        // your cookie settings state. In the future we would refactor the backend preferences to
+        // reflect the only possible states you can be in
+        // (Allow/BlockThirdPartyIncognito/BlockThirdParty/Block), instead of using this
+        // combination of multiple signals.
+        CookieSettingsState state;
+        if (!WebsitePreferenceBridge.isCategoryEnabled(ContentSettingsType.COOKIES)) {
+            state = CookieSettingsState.BLOCK;
+        } else if (PrefServiceBridge.getInstance().getBoolean(Pref.BLOCK_THIRD_PARTY_COOKIES)
+                || PrefServiceBridge.getInstance().getInteger(Pref.COOKIE_CONTROLS_MODE)
+                        == CookieControlsMode.ON) {
+            // Having CookieControlsMode.ON is equivalent to having the BLOCK_THIRD_PARTY_COOKIES
+            // pref set to enabled, because it means third party cookie blocking is always on.
+            state = CookieSettingsState.BLOCK_THIRD_PARTY;
+        } else if (PrefServiceBridge.getInstance().getInteger(Pref.COOKIE_CONTROLS_MODE)
+                == CookieControlsMode.INCOGNITO_ONLY) {
+            state = CookieSettingsState.BLOCK_THIRD_PARTY_INCOGNITO;
+        } else {
+            state = CookieSettingsState.ALLOW;
+        }
+
+        // TODO(crbug.com/1060118): Add managed state for third-party cookie blocking.
+        fourStateCookieToggle.setState(state);
+    }
+
     private void configureTriStateToggle(
             TriStateSiteSettingsPreference triStateToggle, int contentType) {
         triStateToggle.setOnPreferenceChangeListener(this);

chrome/android/java/src/org/chromium/chrome/browser/site_settings/WebsitePreferenceBridge.java

@@ -9,6 +9,7 @@ import androidx.annotation.VisibleForTesting;
 import org.chromium.base.Callback;
 import org.chromium.base.annotations.CalledByNative;
 import org.chromium.base.annotations.NativeMethods;
+import org.chromium.chrome.browser.flags.ChromeFeatureList;
 import org.chromium.chrome.browser.profiles.Profile;
 import org.chromium.components.content_settings.ContentSettingsType;
 import org.chromium.components.location.LocationUtils;
@@ -315,6 +316,15 @@ public class WebsitePreferenceBridge {
         }
     }
 
+    /**
+     * Whether the setting type requires four-state
+     * (Allow/BlockThirdPartyIncognito/BlockThirdParty/Block) setting.
+     */
+    public static boolean requiresFourStateContentSetting(int contentSettingsType) {
+        return contentSettingsType == ContentSettingsType.COOKIES
+                && ChromeFeatureList.isEnabled(ChromeFeatureList.IMPROVED_COOKIE_CONTROLS);
+    }
+
     /**
      * Sets the preferences on whether to enable/disable given setting.
      */

chrome/android/javatests/src/org/chromium/chrome/browser/site_settings/SiteSettingsTest.java

@@ -39,9 +39,11 @@ import org.chromium.chrome.browser.settings.ChromeBaseCheckBoxPreference;
 import org.chromium.chrome.browser.settings.ChromeSwitchPreference;
 import org.chromium.chrome.browser.settings.SettingsActivity;
 import org.chromium.chrome.browser.settings.SettingsLauncher;
+import org.chromium.chrome.browser.site_settings.FourStateCookieSettingsPreference.CookieSettingsState;
 import org.chromium.chrome.test.ChromeActivityTestRule;
 import org.chromium.chrome.test.ChromeJUnit4ClassRunner;
 import org.chromium.chrome.test.util.InfoBarTestAnimationListener;
+import org.chromium.chrome.test.util.browser.Features.DisableFeatures;
 import org.chromium.chrome.test.util.browser.Features.EnableFeatures;
 import org.chromium.chrome.test.util.browser.LocationSettingsTestUtil;
 import org.chromium.components.content_settings.ContentSettingsType;
@@ -63,6 +65,7 @@ import java.util.concurrent.Callable;
 @RetryOnFailure
 @CommandLineFlags.Add({ChromeSwitches.DISABLE_FIRST_RUN_EXPERIENCE,
         ContentSwitches.HOST_RESOLVER_RULES + "=MAP * 127.0.0.1", "ignore-certificate-errors"})
+@EnableFeatures(ChromeFeatureList.IMPROVED_COOKIE_CONTROLS)
 public class SiteSettingsTest {
     @Rule
     public ChromeActivityTestRule<ChromeActivity> mActivityTestRule =
@@ -156,6 +159,29 @@ public class SiteSettingsTest {
     }
 
     private void setCookiesEnabled(final SettingsActivity settingsActivity, final boolean enabled) {
+        TestThreadUtils.runOnUiThreadBlocking(new Runnable() {
+            @Override
+            public void run() {
+                final SingleCategorySettings websitePreferences =
+                        (SingleCategorySettings) settingsActivity.getMainFragment();
+                final FourStateCookieSettingsPreference cookies =
+                        (FourStateCookieSettingsPreference) websitePreferences.findPreference(
+                                SingleCategorySettings.FOUR_STATE_COOKIE_TOGGLE_KEY);
+
+                websitePreferences.onPreferenceChange(
+                        cookies, enabled ? CookieSettingsState.ALLOW : CookieSettingsState.BLOCK);
+                Assert.assertEquals("Cookies should be " + (enabled ? "allowed" : "blocked"),
+                        doesAcceptCookies(), enabled);
+            }
+
+            private boolean doesAcceptCookies() {
+                return WebsitePreferenceBridge.isCategoryEnabled(ContentSettingsType.COOKIES);
+            }
+        });
+    }
+
+    private void setCookiesEnabledOld(
+            final SettingsActivity settingsActivity, final boolean enabled) {
         TestThreadUtils.runOnUiThreadBlocking(new Runnable() {
             @Override
             public void run() {
@@ -184,6 +210,23 @@ public class SiteSettingsTest {
         });
     }
 
+    private void setThirdPartyCookiesEnabledOld(
+            final SettingsActivity settingsActivity, final boolean enabled) {
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            final SingleCategorySettings websitePreferences =
+                    (SingleCategorySettings) settingsActivity.getMainFragment();
+            final ChromeBaseCheckBoxPreference thirdPartyCookies =
+                    (ChromeBaseCheckBoxPreference) websitePreferences.findPreference(
+                            SingleCategorySettings.THIRD_PARTY_COOKIES_TOGGLE_KEY);
+
+            websitePreferences.onPreferenceChange(thirdPartyCookies, enabled);
+            Assert.assertEquals(
+                    "Third-party cookies should be " + (enabled ? "allowed" : "blocked"),
+                    PrefServiceBridge.getInstance().getBoolean(Pref.BLOCK_THIRD_PARTY_COOKIES),
+                    enabled);
+        });
+    }
+
     private void setBlockCookiesSiteException(final SettingsActivity settingsActivity,
             final String hostname, final boolean thirdPartiesOnly) {
         TestThreadUtils.runOnUiThreadBlocking(new Runnable() {
@@ -212,23 +255,6 @@ public class SiteSettingsTest {
         });
     }
 
-    private void setThirdPartyCookiesEnabled(
-            final SettingsActivity settingsActivity, final boolean enabled) {
-        TestThreadUtils.runOnUiThreadBlocking(() -> {
-            final SingleCategorySettings websitePreferences =
-                    (SingleCategorySettings) settingsActivity.getMainFragment();
-            final ChromeBaseCheckBoxPreference thirdPartyCookies =
-                    (ChromeBaseCheckBoxPreference) websitePreferences.findPreference(
-                            SingleCategorySettings.THIRD_PARTY_COOKIES_TOGGLE_KEY);
-
-            websitePreferences.onPreferenceChange(thirdPartyCookies, enabled);
-            Assert.assertEquals(
-                    "Third-party cookies should be " + (enabled ? "allowed" : "blocked"),
-                    PrefServiceBridge.getInstance().getBoolean(Pref.BLOCK_THIRD_PARTY_COOKIES),
-                    enabled);
-        });
-    }
-
     private void setGlobalToggleForCategory(
             final @SiteSettingsCategory.Type int type, final boolean enabled) {
         final SettingsActivity settingsActivity =
@@ -260,6 +286,21 @@ public class SiteSettingsTest {
         settingsActivity.finish();
     }
 
+    private void setFourStateCookieToggle(CookieSettingsState newState) {
+        final SettingsActivity settingsActivity =
+                SiteSettingsTestUtils.startSiteSettingsCategory(SiteSettingsCategory.Type.COOKIES);
+
+        TestThreadUtils.runOnUiThreadBlocking(() -> {
+            SingleCategorySettings preferences =
+                    (SingleCategorySettings) settingsActivity.getMainFragment();
+            FourStateCookieSettingsPreference fourStateCookieToggle =
+                    (FourStateCookieSettingsPreference) preferences.findPreference(
+                            SingleCategorySettings.FOUR_STATE_COOKIE_TOGGLE_KEY);
+            preferences.onPreferenceChange(fourStateCookieToggle, newState);
+        });
+        settingsActivity.finish();
+    }
+
     private void setEnablePopups(final boolean enabled) {
         setGlobalToggleForCategory(SiteSettingsCategory.Type.POPUPS, enabled);
 
@@ -317,13 +358,14 @@ public class SiteSettingsTest {
     @Test
     @SmallTest
     @Feature({"Preferences"})
-    public void testThirdPartyCookieToggleGetsDisabled() {
+    @DisableFeatures(ChromeFeatureList.IMPROVED_COOKIE_CONTROLS)
+    public void testThirdPartyCookieToggleGetsDisabledOld() {
         SettingsActivity settingsActivity =
                 SiteSettingsTestUtils.startSiteSettingsCategory(SiteSettingsCategory.Type.COOKIES);
-        setCookiesEnabled(settingsActivity, true);
-        setThirdPartyCookiesEnabled(settingsActivity, false);
-        setThirdPartyCookiesEnabled(settingsActivity, true);
-        setCookiesEnabled(settingsActivity, false);
+        setCookiesEnabledOld(settingsActivity, true);
+        setThirdPartyCookiesEnabledOld(settingsActivity, false);
+        setThirdPartyCookiesEnabledOld(settingsActivity, true);
+        setCookiesEnabledOld(settingsActivity, false);
         settingsActivity.finish();
     }
 
@@ -538,7 +580,7 @@ public class SiteSettingsTest {
         String[] binaryToggleWithAllowed = new String[] {"binary_toggle", "allowed_group"};
         String[] binaryToggleWithOsWarningExtra =
                 new String[] {"binary_toggle", "os_permissions_warning_extra"};
-        String[] cookie = new String[] {"binary_toggle", "third_party_cookies", "add_exception"};
+        String[] cookie = new String[] {"four_state_cookie_toggle", "add_exception"};
         String[] protectedMedia = new String[] {"tri_state_toggle", "protected_content_learn_more"};
         String[] notifications_enabled;
         String[] notifications_disabled;
@@ -601,6 +643,20 @@ public class SiteSettingsTest {
                 continue;
             }
 
+            // Cookies has a four-state radio preference so it needs to be handled slightly
+            // differently.
+            if (key == SiteSettingsCategory.Type.COOKIES) {
+                setFourStateCookieToggle(CookieSettingsState.ALLOW);
+                checkPreferencesForCategory(key, cookie);
+                setFourStateCookieToggle(CookieSettingsState.BLOCK_THIRD_PARTY_INCOGNITO);
+                checkPreferencesForCategory(key, cookie);
+                setFourStateCookieToggle(CookieSettingsState.BLOCK_THIRD_PARTY);
+                checkPreferencesForCategory(key, cookie);
+                setFourStateCookieToggle(CookieSettingsState.BLOCK);
+                checkPreferencesForCategory(key, cookie);
+                continue;
+            }
+
             Pair<String[], String[]> values = testCases.get(key);
 
             if (key == SiteSettingsCategory.Type.ALL_SITES